Unnecessary fancy indexing in multiple dimensions

Question

Let's say I have a numpy array of size nxmxk and another array B of size nxm that has indices from 1 to k. I want to access each nxm slice of A using the index indicated at that location in B, giving me an array of size nx m.

Edit: this is apparently not what I want! [[I can achieve this using the takefollowing:

A.take(B) ]] end edit

Can this be done using fancy indexing? I would think that it A[B]would give the same result, but the results are in an array of size nxmxmxk (which I really don't understand).

The reason I don't want to use takeis because I want to be able to assign this part, for example

A[B] = 1

The only working solution I have so far is

A.reshape(-1, k)[np.arange(n * m), B.ravel()].reshape(n, m)

but surely there should be an easier way?

Answer 1

Let's pretend that

import numpy as np
np.random.seed(0)

n,m,k = 2,3,5
A = np.arange(n*m*k,0,-1).reshape((n,m,k))
print(A)
# [[[30 29 28 27 26]
#   [25 24 23 22 21]
#   [20 19 18 17 16]]

#  [[15 14 13 12 11]
#   [10  9  8  7  6]
#   [ 5  4  3  2  1]]]

B = np.random.randint(k, size=(n,m))
print(B)
# [[4 0 3]
#  [3 3 1]]

To create this array,

print(A.reshape(-1, k)[np.arange(n * m), B.ravel()])
# [26 25 17 12  7  4]

as an array nxmusing fantasy indexing:

i,j = np.ogrid[0:n, 0:m]
print(A[i, j, B])
# [[26 25 17]
#  [12  7  4]]