Numpy fantazyjne indeksowanie w wielu wymiarach

Question

Załóżmy, że mam tablicę numpy A o rozmiarze n x m x k ​​i inną tablicę B o rozmiarze n x m, która ma indeksy od 1 do k. Chcę uzyskać dostęp do każdego n x m plasterka A za pomocą indeksu podanego w tym miejscu w B, dając mi tablicę o wielkości n x m.

Edytuj: to najwyraźniej nie jest to, czego chcę! [[można to osiągnąć za pomocą take takiego:

A.take(B) ]] end edit

Można to osiągnąć za pomocą ozdobnego indeksowanie? Myślę, że A[B] dałby taki sam wynik, ale to wynik w tablicy o rozmiarach n x m x m x k ​​(których tak naprawdę nie rozumiem).

Powodem nie chcę używać take jest, że chcę, aby móc przypisać tej części czegoś, jak

A[B] = 1

Rozwiązanie działa tylko, że mam tak daleko jest

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

ale na pewno musi być łatwiejszy sposób?

Answer 1

Załóżmy

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

do tworzenia tej tablicy

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

jako nxm tablicy z ozdobnym indeksowania:

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