Kích thước lập chỉ mục kết hợp khó khăn

Question

Sự cố

Tôi có hai mảng gọn gàng, A và indices.

A có kích thước m x ​​n x 10000. indices có kích thước m x ​​n x 5 (đầu ra từ argpartition(A, 5)[:,:,:5]). Tôi muốn lấy một mảng m x n x 5 chứa các phần tử của A tương ứng với indices.

Nỗ lực

indices = np.array([[[5,4,3,2,1],[1,1,1,1,1],[1,1,1,1,1]], 
    [500,400,300,200,100],[100,100,100,100,100],[100,100,100,100,100]]) 
A = np.reshape(range(2 * 3 * 10000), (2,3,10000)) 

A[...,indices] # gives an array of size (2,3,2,3,5). I want a subset of these values 
np.take(A, indices) # shape is right, but it flattens the array first 
np.choose(indices, A) # fails because of shape mismatch. 

Động lực

Tôi đang cố gắng để có được 5 giá trị lớn nhất của A[i,j] cho mỗi i<m, j<n theo thứ tự sắp xếp sử dụng np.argpartition vì các mảng có thể nhận được khá lớn.

Answer 1

Bạn có thể sử dụng advanced-indexing -

m,n = A.shape[:2] 
out = A[np.arange(m)[:,None,None],np.arange(n)[:,None],indices] 

mẫu chạy -

In [330]: A 
Out[330]: 
array([[[38, 21, 61, 74, 35, 29, 44, 46, 43, 38], 
     [22, 44, 89, 48, 97, 75, 50, 16, 28, 78], 
     [72, 90, 48, 88, 64, 30, 62, 89, 46, 20]], 

     [[81, 57, 18, 71, 43, 40, 57, 14, 89, 15], 
     [93, 47, 17, 24, 22, 87, 34, 29, 66, 20], 
     [95, 27, 76, 85, 52, 89, 69, 92, 14, 13]]]) 

In [331]: indices 
Out[331]: 
array([[[7, 8, 1], 
     [7, 4, 7], 
     [4, 8, 4]], 

     [[0, 7, 4], 
     [5, 3, 1], 
     [1, 4, 0]]]) 

In [332]: m,n = A.shape[:2] 

In [333]: A[np.arange(m)[:,None,None],np.arange(n)[:,None],indices] 
Out[333]: 
array([[[46, 43, 21], 
     [16, 97, 16], 
     [64, 46, 64]], 

     [[81, 14, 43], 
     [87, 24, 47], 
     [27, 52, 95]]]) 

Để nhận được những chỉ số tương ứng với tối đa 5 yếu tố dọc theo trục qua, chúng tôi sẽ sử dụng argpartition, như vậy -

indices = np.argpartition(-A,5,axis=-1)[...,:5] 

Để giữ trật tự từ cao nhất đến thấp nhất, chúng tôi e range(5) thay vì 5.

Answer 2

Đối với hậu thế, sau đây sử dụng câu trả lời Divakar để hoàn thành mục tiêu ban đầu, ví dụ: trả lại 5 giá trị hàng đầu cho tất cả i<m, j<n theo thứ tự sắp xếp:

m, n = np.shape(A)[:2] 

# get the largest 5 indices for all m, n 
top_unsorted_indices = np.argpartition(A, -5, axis=2)[...,-5:] 

# get the values corresponding to top_unsorted_indices 
top_values = A[np.arange(m)[:,None,None], np.arange(n)[:,None], top_unsorted_indices] 

# sort the top 5 values 
top_sorted_indices = top_unsorted_indices[np.arange(m)[:,None,None], np.arange(n)[:,None], np.argsort(-top_values)]