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)
有没有更向量化的方法呢?
问题答案:
我们可以将coords
as的每一行拆开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)