Question

我有这个2d的零z数组和这个1d的起点starts数组。另外，我有一个offsets

的1d数组

z = z = np.zeros(35, dtype='i').reshape(5, 7)
starts = np.array([1, 5, 3, 0, 3])
offsets = np.arange(5) + 1

我想在这里对这个小循环进行矢量化，但我似乎无法做到这一点。

for i in range(z.shape[0]):
    z[i, starts[i]:] += offsets[i]

此示例中的结果应如下所示：

z
array([[0, 1, 1, 1, 1, 1, 1],
       [0, 0, 0, 0, 0, 2, 2],
       [0, 0, 0, 3, 3, 3, 3],
       [4, 4, 4, 4, 4, 4, 4],
       [0, 0, 0, 5, 5, 5, 5]])

Answer 1

我们可以使用一些masking和NumPy broadcasting -

mask = starts[:,None] <= np.arange(z.shape[1])
z[mask] = np.repeat(offsets, mask.sum(1))

我们可以发挥broadcasted multiplication的技巧来获得最终输出 -

z = offsets[:,None] * mask

其他方法是将值从z分配到offsets，然后屏蔽mask的其余部分，就像这样 -

z[:] = offsets[:,None]
z[~mask] = 0

其他方式是将offsets的复制版本作为起始z然后屏蔽 -

z = np.repeat(offsets,z.shape[1]).reshape(z.shape[0],-1)
z[~mask] = 0

当然，我们需要先前的形状参数。

如果z未初始化为zeros数组，那么前面提到的解决方案中只有一个适用，而且需要使用+=进行更新，就像这样 -

z[mask] += np.repeat(offsets, mask.sum(1))

numpy索引：从不同的位置开始，向行的部分添加向量

1 个答案: