Tricky numpy indexing

user_378147 0
Viewed 46

I have values in a 2D array called a. I have a second array, C, in which I have indices which refers to cells in a. This array's dimensions are therefore (M, N, 2).

I want to construct the output of dimensions (M, N) which is made of the values of a from the indices stored in C.

This is a pythonic way to do so:

a = np.arange(20).reshape(-1, 5)
print(a)
# [[ 0  1  2  3  4]
#  [ 5  6  7  8  9]
#  [10 11 12 13 14]
#  [15 16 17 18 19]]

C = np.array([
    [[2, 0], [3, 1]],
    [[0, 0], [1, 4]],
    [[2, 2], [1, 1]],
    [[0, 0], [1, 1]],
])


output = np.array([[a[cell[0], cell[1]] for cell in row] for row in C])
print(output)

# [[10 16]
#  [ 0  9]
#  [12  6]
#  [ 0  6]]

I guess there is a more numpy-ic way which would be much more elegant but, more importantly, much more efficient to do the same thing.

Any ideas?

Cheers

numpy arrays python
1 Answers

@Divakar: This is indeed the solution. I'm so amazed... Didn't event know that was valid code. Thanks a lot.

>>> output = a[C[..., 0], C[..., 1]]
>>> output
array([[10, 16],
       [ 0,  9],
       [12,  6],
       [ 0,  6]])
user_3781470
Related