Memilih baris dan kolom tertentu dari array NumPy

Aku sudah gila mencoba mencari tahu hal bodoh apa yang aku lakukan salah di sini.

Saya menggunakan NumPy, dan saya memiliki indeks baris tertentu dan indeks kolom tertentu yang ingin saya pilih. Inilah inti dari masalah saya:

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

Mengapa ini terjadi? Tentunya saya harus bisa memilih baris 1, 2, dan 4, dan kolom 1 dan 3? Hasil yang saya harapkan adalah:

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

Pengindeksan mewah mengharuskan Anda memberikan semua indeks untuk setiap dimensi. Anda memberikan 3 indeks untuk yang pertama, dan hanya 2 untuk yang kedua, oleh karena itu terjadi kesalahan. Anda ingin melakukan sesuatu seperti ini:

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

Itu tentu saja merepotkan untuk menulis, jadi Anda dapat membiarkan penyiaran membantu Anda:

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

Ini jauh lebih mudah dilakukan jika Anda mengindeks dengan array, bukan daftar:

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

Seperti Toan menyarankan, hack sederhana akan cukup pilih baris pertama, dan kemudian pilih kolom atas yang .

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

[Sunting] Metode bawaan: np.ix_

Saya baru-baru menemukan bahwa numpy memberi Anda built-in satu-kapal untuk melakukan persis apa yang @Jaime disarankan, tetapi tanpa harus menggunakan penyiaran sintaks (yang menderita dari kurangnya dibaca). Dari dokumen:

Menggunakan ix_ satu dapat dengan cepat membuat array indeks yang akan mengindeks produk silang. a[np.ix_([1,3],[2,5])]mengembalikan array [[a[1,2] a[1,5]], [a[3,2] a[3,5]]].

Jadi Anda menggunakannya seperti ini:

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

Dan cara kerjanya adalah menangani penyelarasan array seperti yang disarankan Jaime, sehingga penyiaran terjadi dengan benar:

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

Juga, seperti yang dikatakan MikeC dalam sebuah komentar, np.ix_memiliki keuntungan mengembalikan tampilan, yang tidak didapatkan pada jawaban (pra-edit) pertama saya. Ini berarti Anda sekarang dapat menetapkan ke array yang diindeks:

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

MENGGUNAKAN:

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

ATAU:

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

Menggunakan np.ix_adalah cara paling mudah untuk melakukannya (seperti yang dijawab oleh orang lain), tetapi berikut ini cara lain yang menarik untuk melakukannya:

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

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

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