根据部分总和添加填充

Question

我有四个给定的变量：

• 组大小
• 总和
• 部分总和
• 1-D张量 当组中的总和达到部分总和时，我想添加零。例如：

groupsize = 4
totalgroups = 3
partialsum = 15
d1tensor = torch.tensor([ 3, 12,  5,  5,  5,  4, 11])

预期的结果是： [3，12，0，0，5，5，5，5，0，4，11，0，0]

我不知道如何在纯pytorch中实现这一目标。在Python中，这将是这样的：

target = [0]*(groupsize*totalgroups)

cursor = 0
current_count = 0
d1tensor = [ 3, 12,  5,  5,  5,  4, 11]

for idx, ele in enumerate(target):
  subgroup_start = (idx//groupsize) *groupsize
  subgroup_end = subgroup_start + groupsize 
  if sum(target[subgroup_start:subgroup_end]) < partialsum: 
     target[idx] = d1tensor[cursor]
     cursor +=1

有人可以帮我吗？我已经搜索过它，但找不到任何东西。

Answer 1

一些逻辑，numpy和列表综合在这里就足够了。
我会逐步将其分解，您可以使其更苗条，更漂亮：

import numpy as np

my_val = 15
block_size = 4
total_groups = 3
d1 = [3, 12,  5,  5,  5,  4, 11]
d2 = np.cumsum(d1)
d3 = d2 % my_val == 0 #find where sum of elements is 15 or multiple
split_points= [i+1 for i, x in enumerate(d3) if x] # find index where cumsum == my_val
#### Option 1
split_array = np.split(d1, split_points, axis=0)
padded_arrays = [np.pad(array, (0, block_size - len(array)), mode='constant') for array in split_array] #pad arrays
padded_d1 = np.concatenate(padded_arrays[:total_groups]) #put them together, discard extra group if present
#### Option 2
split_points = [el for el in split_points if el <len(d1)] #make sure we are not splitting on the last element of d1
split_array = np.split(d1, split_points, axis=0)
padded_arrays = [np.pad(array, (0, block_size - len(array)), mode='constant') for array in split_array] #pad arrays
padded_d1 = np.concatenate(padded_arrays)