Fill certain indices of a 3d numpy array

Question

I have a list of indices likes this:

selected_coords = [[1, 8, 30], [15, 4, 6] ,...]

And a list of values like this:

differences = [1, 5, 8, 2, ...]

Both have 500 entries. Now I want to fill a 3d numpy array with these values on the right index. What I tried to do is the following:

brain_map = np.zeros(shape=(48,60,22))

for i, index in enumerate(selected_coords):
    ind = list(map(int, index))
    brain_map[ind] = differences[i]

If I print the index and the value in this loop I get the right format, but if I print the matrix after the loop it seems like the values have been put in there multiple times instead of only on the specified indices. What am I doing wrong?

Answer 1

You should avoid looping over numpy arrays whenever possible, otherwise you're losing performance. You can make use of advanced ("fancy") indexing to index a subset of elements at specific indices. This would work like so:

brain_map[ind_x, ind_y, ind_z] = vals

where ind_x, ind_y, ind_z and vals are all 1d array-likes of the same length. What you have is essentially the transpose of your index arrays:

brain_map[tuple(zip(*selected_coords))] = differences

The zip(*) trick essentially transposes your list of lists, which can then be passed as a tuple for indexing. For instance:

>>> import numpy as np
>>> M = np.random.rand(2, 3, 4)
>>> coords = [[0, 1, 2], [1, 2, 3]]
>>> tuple(zip(*coords))
((0, 1), (1, 2), (2, 3))
>>> M[tuple(zip(*coords))]
array([ 0.12299864,  0.76461622])
>>> M[0, 1, 2], M[1, 2, 3]
(0.12299863762892316, 0.76461622348724623)