argsort cho một ndarray đa chiều

Question

Tôi đang cố gắng để có được các chỉ số sắp xếp một mảng đa chiều theo trục cuối cùng, ví dụ:

>>> a = np.array([[3,1,2],[8,9,2]]) 

Và tôi muốn chỉ số i như vậy mà,

>>> a[i] 
array([[1, 2, 3], 
     [2, 8, 9]]) 

Dựa trên các tài liệu của numpy.argsort Tôi nghĩ rằng nó nên làm điều này, nhưng tôi nhận được lỗi:

>>> a[np.argsort(a)] 
IndexError: index 2 is out of bounds for axis 0 with size 2 

---

Chỉnh sửa: Tôi cần phải sắp xếp lại các mảng khác có cùng hình dạng (ví dụ: mảng b mà a.shape == b.shape) trong cùng một cách ... do đó

>>> b = np.array([[0,5,4],[3,9,1]]) 
>>> b[i] 
array([[5,4,0], 
     [9,3,1]]) 

Answer 1

Giải pháp:

>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)] 
array([[1, 2, 3], 
     [2, 8, 9]]) 

---

Bạn đã làm đúng, mặc dù tôi sẽ không mô tả nó như gian lận việc lập chỉ mục.

Có lẽ điều này sẽ giúp làm cho nó rõ ràng hơn:

In [544]: i=np.argsort(a,axis=1) 

In [545]: i 
Out[545]: 
array([[1, 2, 0], 
     [2, 0, 1]]) 

i là thứ tự mà chúng ta muốn, cho mỗi hàng. Đó là:

In [546]: a[0, i[0,:]] 
Out[546]: array([1, 2, 3]) 

In [547]: a[1, i[1,:]] 
Out[547]: array([2, 8, 9]) 

Để thực hiện cả hai bước lập chỉ mục cùng một lúc, chúng tôi phải sử dụng chỉ mục 'cột' cho thứ nguyên thứ nhất.

In [548]: a[[[0],[1]],i] 
Out[548]: 
array([[1, 2, 3], 
     [2, 8, 9]]) 

Một mảng có thể được ghép nối với i là:

In [560]: j=np.array([[0,0,0],[1,1,1]]) 

In [561]: j 
Out[561]: 
array([[0, 0, 0], 
     [1, 1, 1]]) 

In [562]: a[j,i] 
Out[562]: 
array([[1, 2, 3], 
     [2, 8, 9]]) 

Nếu i xác định các cột cho mỗi phần tử, sau đó j xác định hàng cho mỗi yếu tố. Mảng cột [[0],[1]] hoạt động tốt bởi vì nó có thể được phát sóng với i.

tôi nghĩ về

np.array([[0], 
      [1]]) 

là 'tay ngắn' cho j. Họ cùng nhau xác định hàng và cột nguồn của mỗi phần tử của mảng mới. Họ làm việc cùng nhau, không phải tuần tự.

Việc lập bản đồ đầy đủ từ a đến các mảng mới là:

[a[0,1] a[0,2] a[0,0] 
a[1,2] a[1,0] a[1,1]] 

---

def foo(a): 
    i = np.argsort(a, axis=1) 
    return (np.arange(a.shape[0])[:,None], i) 

In [61]: foo(a) 
Out[61]: 
(array([[0], 
     [1]]), array([[1, 2, 0], 
     [2, 0, 1]], dtype=int32)) 
In [62]: a[foo(a)] 
Out[62]: 
array([[1, 2, 3], 
     [2, 8, 9]]) 

Answer 2

tôi thấy the answer here, với một người nào đó có cùng một vấn đề. Họ chủ yếu chỉ là lừa dối việc lập chỉ mục để làm việc đúng cách ...

>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)] 
array([[1, 2, 3], 
     [2, 8, 9]]) 

Answer 3

Bạn cũng có thể sử dụng linear indexing, mà có thể được tốt hơn với hiệu suất, như vậy -

M,N = a.shape 
out = b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)] 

Vì vậy, a.argsort(1)+(np.arange(M)[:,None]*N) về cơ bản là các chỉ số tuyến tính được sử dụng để lập bản đồ b để có được kết quả được sắp xếp mong muốn cho b. Chỉ số tuyến tính tương tự cũng có thể được sử dụng trên a để nhận đầu ra được sắp xếp cho a.

mẫu chạy - kiểm tra

In [23]: a = np.array([[3,1,2],[8,9,2]]) 

In [24]: b = np.array([[0,5,4],[3,9,1]]) 

In [25]: M,N = a.shape 

In [26]: b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)] 
Out[26]: 
array([[5, 4, 0], 
     [1, 3, 9]]) 

Rumtime -

In [27]: a = np.random.rand(1000,1000) 

In [28]: b = np.random.rand(1000,1000) 

In [29]: M,N = a.shape 

In [30]: %timeit b[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)] 
10 loops, best of 3: 133 ms per loop 

In [31]: %timeit b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)] 
10 loops, best of 3: 96.7 ms per loop