numpy:将值分配给带有索引列表的2d数组


问题内容

我有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)

有没有更向量化的方法呢?


问题答案:

我们可以将coordsas的每一行拆开img,然后将col索引拆包以建立索引,然后进行分配。

现在,由于该问题被标记为: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)