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

Question

我有这个二维零数组 z 和这个一维起点数组 starts。另外，我有一个 offsets

的一维数组

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

我想在这里对这个小 for 循环进行矢量化，但我似乎做不到。

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 and NumPy broadcasting -

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

我们可以玩broadcasted multiplication的把戏来得到最终的输出-

z = offsets[:,None] * mask

其他方法是将值从 offsets 分配到 z，然后屏蔽掉 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))