Numpy:使用索引列表为2d数组赋值

时间:2017-10-30 12:42:23

标签: python python-3.x numpy vectorization

我有2d numpy数组(想想灰度图像)。我想将特定值分配给此数组的坐标列表,例如:

img = np.zeros((5, 5))
coords = np.array([[0, 1], [1, 2], [2, 3], [3, 4]]) 

def bad_use_of_numpy(img, coords):
    for i, coord in enumerate(coords):
        img[coord[0], coord[1]] = 255

    return img

bad_use_of_numpy(img, coords)

这很有效,但我觉得我可以利用numpy功能来加快速度。我之后也可能会遇到类似以下的用例:

img = np.zeros((5, 5))
coords = np.array([[0, 1], [1, 2], [2, 3], [3, 4]])
vals = np.array([1, 2, 3, 4])

def bad_use_of_numpy(img, coords, vals):
    for coord in coords:
        img[coord[0], coord[1]] = vals[i]

    return img

 bad_use_of_numpy(img, coords, vals)

有更多的矢量化方式吗?

1 个答案:

答案 0 :(得分:6)

我们可以将coords的每一行解包为row,col索引,以便索引到img然后分配。

现在,由于问题被标记为:Python 3.x,我们可以使用[*coords.T]解压缩然后分配 -

img[[*coords.T]] = 255

通常,我们可以使用tuple解包 -

img[tuple(coords.T)] = 255

我们还可以计算线性指数,然后使用np.put -

进行分配
np.put(img, np.ravel_multi_index(coords.T, img.shape), 255)