Numpy: assigning values to 2d array with list of i

2020-02-14 03:47发布

问题:

I have 2d numpy array (think greyscale image). I want to assign certain value to a list of coordinates to this array, such that:

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)

This works, but I feel like I can take advantage of numpy functionality to make it faster. I also might have a use case later to to something like following:

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)

Is there a more vectorized way of doing that?

回答1:

We can unpack each row of coords as row, col indices for indexing into img and then assign.

Now, since the question is tagged : Python 3.x, on it we can simply unpack with [*coords.T] and then assign -

img[[*coords.T]] = 255

Generically, we can use tuple to unpack -

img[tuple(coords.T)] = 255

We can also compute the linear indices and then assign with np.put -

np.put(img, np.ravel_multi_index(coords.T, img.shape), 255)