Chọn các hàng và cột cụ thể từ mảng NumPy

Tôi đã phát điên khi cố gắng tìm ra điều ngu ngốc mà tôi đang làm sai ở đây.

Tôi đang sử dụng NumPy và tôi có chỉ số hàng và chỉ số cột cụ thể mà tôi muốn chọn. Đây là ý chính của vấn đề của tôi:

import numpy as np

a = np.arange(20).reshape((5,4))
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11],
#        [12, 13, 14, 15],
#        [16, 17, 18, 19]])

# If I select certain rows, it works
print a[[0, 1, 3], :]
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [12, 13, 14, 15]])

# If I select certain rows and a single column, it works
print a[[0, 1, 3], 2]
# array([ 2,  6, 14])

# But if I select certain rows AND certain columns, it fails
print a[[0,1,3], [0,2]]
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# ValueError: shape mismatch: objects cannot be broadcast to a single shape

Tại sao chuyện này đang xảy ra? Chắc chắn tôi sẽ có thể chọn các hàng 1, 2 và 4, và các cột 1 và 3? Kết quả tôi mong đợi là:

a[[0,1,3], [0,2]] => [[0,  2],
                      [4,  6],
                      [12, 14]]

Lập chỉ mục lạ mắt yêu cầu bạn cung cấp tất cả các chỉ số cho từng thứ nguyên. Bạn đang cung cấp 3 chỉ số cho chỉ số đầu tiên và chỉ 2 chỉ số cho chỉ số thứ hai, do đó lỗi. Bạn muốn làm điều gì đó như sau:

>>> a[[[0, 0], [1, 1], [3, 3]], [[0,2], [0,2], [0, 2]]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Đó dĩ nhiên là một điều khó khăn khi viết, vì vậy bạn có thể để việc phát sóng giúp bạn:

>>> a[[[0], [1], [3]], [0, 2]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Điều này đơn giản hơn nhiều để thực hiện nếu bạn lập chỉ mục với mảng, không phải danh sách:

>>> row_idx = np.array([0, 1, 3])
>>> col_idx = np.array([0, 2])
>>> a[row_idx[:, None], col_idx]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Như Toàn gợi ý, một cách hack đơn giản là chỉ cần chọn các hàng trước, sau đó chọn các cột trên đó .

>>> a[[0,1,3], :]            # Returns the rows you want
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [12, 13, 14, 15]])
>>> a[[0,1,3], :][:, [0,2]]  # Selects the columns you want as well
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

[Chỉnh sửa] Phương thức tích hợp: np.ix_

Gần đây tôi đã phát hiện ra rằng numpy cung cấp cho bạn một lớp lót tích hợp để thực hiện chính xác những gì @Jaime đề xuất, nhưng không cần phải sử dụng cú pháp phát sóng (thiếu khả năng đọc). Từ các tài liệu:

Sử dụng ix_ người ta có thể nhanh chóng tạo các mảng chỉ mục sẽ lập chỉ mục cho sản phẩm chéo. a[np.ix_([1,3],[2,5])]trả về mảng [[a[1,2] a[1,5]], [a[3,2] a[3,5]]].

Vì vậy, bạn sử dụng nó như thế này:

>>> a = np.arange(20).reshape((5,4))
>>> a[np.ix_([0,1,3], [0,2])]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Và cách nó hoạt động là nó sẽ chăm sóc sắp xếp các mảng theo cách mà Jaime đã đề xuất, để việc phát sóng diễn ra đúng cách:

>>> np.ix_([0,1,3], [0,2])
(array([[0],
        [1],
        [3]]), array([[0, 2]]))

Ngoài ra, như MikeC nói trong một bình luận, np.ix_có lợi thế là trả về một lượt xem, điều mà câu trả lời đầu tiên (trước khi chỉnh sửa) của tôi không có. Điều này có nghĩa là bây giờ bạn có thể gán cho mảng được lập chỉ mục:

>>> a[np.ix_([0,1,3], [0,2])] = -1
>>> a    
array([[-1,  1, -1,  3],
       [-1,  5, -1,  7],
       [ 8,  9, 10, 11],
       [-1, 13, -1, 15],
       [16, 17, 18, 19]])

SỬ DỤNG:

 >>> a[[0,1,3]][:,[0,2]]
array([[ 0,  2],
   [ 4,  6],
   [12, 14]])

HOẶC LÀ:

>>> a[[0,1,3],::2]
array([[ 0,  2],
   [ 4,  6],
   [12, 14]])

Sử dụng np.ix_là cách thuận tiện nhất để làm điều đó (như những người khác đã trả lời), nhưng đây là một cách thú vị khác để làm điều đó:

>>> rows = [0, 1, 3]
>>> cols = [0, 2]

>>> a[rows].T[cols].T

array([[ 0,  2],
       [ 4,  6],
       [12, 14]])