多維度的numpy花式索引

Question

假設我有一個大小爲n x m x k的numpy數組A和另一個大小爲n x m的數組B，它的索引從1到k。 我想訪問每個n×m切片A使用在這個地方給出的索引B， 給我一個大小爲n x m的數組。

編輯：這顯然不是我想要的！ [我可以實現使用take這個是這樣的：

A.take(B) ]編輯完

可以這樣使用花哨的索引來實現？ 我本來以爲A[B]會給出相同的結果，但結果 在一個大小爲n x m x m x k（我真的不明白）的數組中。

的原因，我不希望使用take是，我希望能夠在這部分上指定的東西，像

A[B] = 1

唯一的工作解決方案，我至今是

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

但肯定有一個更簡單的方法？

Answer 1

假設

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]] 

要創建此陣列，

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

作爲nxm陣列使用花式索引：

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