Memilih baris dan kolom tertentu dari array NumPy

97

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]]
Mike C.
sumber
Tagged numpy-slicing untuk meningkatkan kemudahan ditemukan. (Juga istilah 'slice' dan 'slicing' tidak muncul dalam teks biasa, kita dapat menggunakan beberapa duplikat dengan istilah tersebut ditutup dengan ini)
smci
Ini adalah duplikat dari stackoverflow.com/questions/19161512/numpy-extract-submatrix
David John Coleman II

Jawaban:

86

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]])
Jaime
sumber
4
Terima kasih, saya tidak tahu Anda bisa melakukan ini! Penyiaran itu aneh dan luar biasa ... Setelah dua tahun numpy, saya masih terbiasa dengannya.
Praveen
2
Terima kasih! Sementara jawaban lain menjawab pertanyaan saya dengan benar dalam hal mengembalikan matriks yang dipilih, jawaban ini menjawabnya sambil juga membahas masalah penugasan (cara menetapkan [[0,1,3], [0,2]] = 0 , sebagai contoh).
Mike C
1
@Jaime - Baru kemarin saya menemukan one-liner built-in untuk melakukan persis trik penyiaran yang Anda sarankan: np.ix_
Praveen
1
Bisakah seseorang memberikan penjelasan mengapa sintaksnya bekerja seperti ini? Apa alasannya berfungsi untuk kedua contoh pertama tetapi tidak untuk yang ketiga. Dan juga, bagaimana merangkum indeks yang diinginkan dalam daftar mereka sendiri menyelesaikan masalah ini? Terima kasih
Aetos
2
Mengapa baris perlu disarangkan dan kolom tidak?
AturSams
86

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]])
Praveen
sumber
4
Dalam beberapa tes, saya juga menemukan np.ix_lebih cepat daripada metode memilih kolom pertama dan kemudian baris (biasanya sekitar 2x lebih cepat pada pengujian saya terhadap array persegi berukuran 1K-10K di mana Anda mengindeks ulang semua baris dan kolom).
Nathan
7

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]])
Toan Nguyen
sumber
10
Meskipun ini benar, Anda harus mempertimbangkan untuk memposting sedikit informasi lebih lanjut yang menjelaskan mengapa itu benar.
ebarr
2

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]])
Andreas K.
sumber