向外压缩numpy码

Question

我试着把下面的代码矢量化-

for i in range(imL.shape[0]):
    for j in range(imL.shape[1]):
        if j - disp[i][j] >= 0:
            imR[i, j - disp[i][j], :] = imL[i, j, :]

基本上，我想这么做-

A[i, j - B[i, j], k] = C[i, j, k] for all (i, j, k)

我查过布尔索引，想出了这个-

tmp = np.arange(imR.shape[1])
tmp = np.repeat(tmp[np.newaxis, :], imR.shape[0], axis=0)
tmp = np.repeat(tmp[:, :, np.newaxis], imR.shape[2], axis=2)

imR[(tmp - disp) >= 0] = imL

但是，这会抛出错误-

TypeError: NumPy boolean array indexing assignment requires a 0 or 1-dimensional input, input has 3 dimensions

我相信整数索引是解决方案，然而，我不能想出任何东西。有没有一种有效的向量化方法？

注: imL和imR为NxMx3。disp是一个N ndarray。

Answer 1

可以使用np.indices简化临时索引数组的创建。

i_indices, j_indices = np.indices(disp.shape)

然后可以使用自定义公式创建更新的索引数组：

# j - B[i, k]
new_j_ind = j_indices - disp

在以下条件下，用原始的j值替换更新的j值：

# if j - disp[i][j] < 0, keep original j index
new_j_ind[new_j_ind < 0] = j_indices[new_j_ind < 0]

然后写数组。

# A[i, j - B[i, j] if j - B[i, j] > 0 else j, k] = C[i, j, k] for all (i, j, k)
imR[i_indices, new_j_ind] = imL

但是要小心:如果[i, j - B[i, j]]对任何(i，k)都映射到相同的坐标，这是没有很好定义的。在for循环中，它定义得很好(最后写的值获胜)，在numpy向量化代码中则不是。