NumPy chọn chỉ mục cột cụ thể cho mỗi hàng bằng cách sử dụng danh sách chỉ mục

Tôi đang đấu tranh để chọn các cột cụ thể trên mỗi hàng của NumPyma trận.

Giả sử tôi có ma trận sau đây mà tôi sẽ gọi X:

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]

Tôi cũng có một listsố chỉ mục cột trên mỗi hàng mà tôi sẽ gọi Y:

[1, 0, 2]

Tôi cần lấy các giá trị:

[2]
[4]
[9]

Thay vì a listvới các chỉ mục Y, tôi cũng có thể tạo một ma trận có hình dạng giống như Xtrong đó mọi cột là a bool/ inttrong phạm vi giá trị 0-1, cho biết đây có phải là cột bắt buộc hay không.

[0, 1, 0]
[1, 0, 0]
[0, 0, 1]

Tôi biết điều này có thể được thực hiện bằng cách lặp lại mảng và chọn các giá trị cột mà tôi cần. Tuy nhiên, điều này sẽ được thực thi thường xuyên trên các mảng dữ liệu lớn và đó là lý do tại sao nó phải chạy nhanh nhất có thể.

Vì vậy, tôi đã tự hỏi liệu có một giải pháp tốt hơn không?

Cảm ơn bạn.

Nếu bạn có một mảng boolean, bạn có thể thực hiện lựa chọn trực tiếp dựa trên như vậy:

>>> a = np.array([True, True, True, False, False])
>>> b = np.array([1,2,3,4,5])
>>> b[a]
array([1, 2, 3])

Để đi cùng với ví dụ ban đầu của bạn, bạn có thể làm như sau:

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> b = np.array([[False,True,False],[True,False,False],[False,False,True]])
>>> a[b]
array([2, 4, 9])

Bạn cũng có thể thêm arangevà thực hiện lựa chọn trực tiếp trên đó, mặc dù tùy thuộc vào cách bạn đang tạo mảng boolean của mình và mã của bạn trông giống như YMMV.

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> a[np.arange(len(a)), [1,0,2]]
array([2, 4, 9])

Hy vọng điều đó sẽ hữu ích, hãy cho tôi biết nếu bạn có thêm bất kỳ câu hỏi nào.

Bạn có thể làm điều gì đó như sau:

In [7]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [8]: lst = [1, 0, 2]

In [9]: a[np.arange(len(a)), lst]
Out[9]: array([2, 4, 9])

Thông tin thêm về lập chỉ mục mảng đa chiều: http://docs.scipy.org/doc/numpy/user/basics.indexing.html#indexing-multi-dimensional-arrays

Một cách đơn giản có thể giống như sau:

In [1]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [2]: y = [1, 0, 2]  #list of indices we want to select from matrix 'a'

range(a.shape[0]) sẽ trở lại array([0, 1, 2])

In [3]: a[range(a.shape[0]), y] #we're selecting y indices from every row
Out[3]: array([2, 4, 9])

Các numpyphiên bản gần đây đã thêm một take_along_axis(và put_along_axis) thực hiện việc lập chỉ mục này một cách rõ ràng.

In [101]: a = np.arange(1,10).reshape(3,3)                                                             
In [102]: b = np.array([1,0,2])                                                                        
In [103]: np.take_along_axis(a, b[:,None], axis=1)                                                     
Out[103]: 
array([[2],
       [4],
       [9]])

Nó hoạt động theo cách tương tự như:

In [104]: a[np.arange(3), b]                                                                           
Out[104]: array([2, 4, 9])

nhưng với cách xử lý trục khác nhau. Nó đặc biệt nhằm vào việc áp dụng các kết quả của argsortvà argmax.

Bạn có thể làm điều đó bằng cách sử dụng trình lặp. Như thế này:

np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)

Thời gian:

N = 1000
X = np.zeros(shape=(N, N))
Y = np.arange(N)

#@Aशwini चhaudhary
%timeit X[np.arange(len(X)), Y]
10000 loops, best of 3: 30.7 us per loop

#mine
%timeit np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)
1000 loops, best of 3: 1.15 ms per loop

#mine
%timeit np.diag(X.T[Y])
10 loops, best of 3: 20.8 ms per loop

Một cách thông minh khác là chuyển đổi mảng đầu tiên và lập chỉ mục nó sau đó. Cuối cùng, lấy đường chéo, nó luôn là câu trả lời đúng.

X = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
Y = np.array([1, 0, 2, 2])

np.diag(X.T[Y])

Từng bước một:

Mảng ban đầu:

>>> X
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

>>> Y
array([1, 0, 2, 2])

Transpose để có thể lập chỉ mục đúng.

>>> X.T
array([[ 1,  4,  7, 10],
       [ 2,  5,  8, 11],
       [ 3,  6,  9, 12]])

Nhận các hàng theo thứ tự Y.

>>> X.T[Y]
array([[ 2,  5,  8, 11],
       [ 1,  4,  7, 10],
       [ 3,  6,  9, 12],
       [ 3,  6,  9, 12]])

Đường chéo bây giờ sẽ trở nên rõ ràng.

>>> np.diag(X.T[Y])
array([ 2,  4,  9, 12]