'What is the best way in numpy to use a vector of positions as indexes in a an nd array?

I have a numpy array of agents positions:

positions = np.array([[row_0, col_0],
                      [row_1, col_1],
                      [row_2, col_2]])

I have a 2d grid, grid = np.zeros((N, N)). Now, let's say I want to mark all the cells where the agents are with True. What I have been doing is this:

grid[positions[:, 0], positions[:, 1]] = True

However, this is wordy and not scalable. Let's say now the positions are in 3d space. Now we have to modify the code like so:

grid[positions[:, 0], positions[:, 1], positions[:, 2]] = True

There is another way to solve the scalability problem, which is to do this:

grid[tuple(positions.T)] = True

However, it a involves a function call. Additionally, let's say there is not one, but two grids, such as:

grid = np.zeros((2, N, N))

and I wanted to mark only the first grid with True, then the method above does not work unless I do:

grid[0, (*tuple(positions.T))] = True

which does not look readable, and I wonder if there is a better way to do this. Any help is appreciated!

EDIT: I have been looking into this, and it finally hit me that adding a comma , in the parenthesis will allow tuple unpacking without explicitly calling the tuple constructor, like so:

grid[(*positions.T,)] = True

For the case of modifying the grid at index n of many grids combined at axis 0, the command becomes:

grid[(n, *positions.T,)] = True


Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source