Numpy ưa thích lập chỉ mục trong nhiều kích thước

Question

Giả sử tôi có mảng numpy A có kích thước n x m x k ​​và một mảng B có kích thước n x m có chỉ mục từ 1 đến k. Tôi muốn truy cập từng lát n x m của A bằng chỉ mục được cung cấp tại địa điểm này trong B, cho tôi một mảng có kích thước n x m.

Chỉnh sửa: đó rõ ràng không phải là những gì tôi muốn! [[tôi có thể đạt được điều này bằng take như thế này:

A.take(B) ]] cuối chỉnh sửa

thể này đạt được sử dụng lập chỉ mục ưa thích? Tôi có thể nghĩ rằng A[B] sẽ cho kết quả tương tự, nhưng kết quả là trong một mảng có kích thước n x m x m x k ​​(mà tôi không thực sự hiểu).

Lý do tôi không muốn sử dụng take là tôi muốn để có thể gán phần một cái gì đó này, giống như

A[B] = 1

Giải pháp duy nhất làm việc mà tôi có cho đến nay là

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

nhưng chắc chắn phải có cách dễ dàng hơn?

Answer 1

Giả sử

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

Để tạo mảng này,

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

như một nxm mảng sử dụng lập chỉ mục ưa thích:

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