Выбор определенных строк и столбцов из массива NumPy
Я схожу с ума, пытаясь понять, какую глупость я делаю здесь неправильно.
Я использую NumPy, и у меня есть конкретные индексы строк и конкретные индексы столбцов, из которых я хочу выбрать. Вот суть моей проблемы:
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
Почему это происходит? Конечно, я должен иметь возможность выбрать 1-й, 2-й и 4-й строки, а также 1-й и 3-й столбцы? Результат, который я ожидаю, таков:
a[[0,1,3], [0,2]] => [[0, 2],
[4, 6],
[12, 14]]
3 ответа
Необычное индексирование требует, чтобы вы предоставили все индексы для каждого измерения. Вы предоставляете 3 индекса для первого и только 2 для второго, следовательно, ошибка. Вы хотите сделать что-то вроде этого:
>>> a[[[0, 0], [1, 1], [3, 3]], [[0,2], [0,2], [0, 2]]]
array([[ 0, 2],
[ 4, 6],
[12, 14]])
Это, конечно, боль писать, поэтому вы можете позволить вещанию помочь вам:
>>> a[[[0], [1], [3]], [0, 2]]
array([[ 0, 2],
[ 4, 6],
[12, 14]])
Это гораздо проще сделать, если вы индексируете с помощью массивов, а не списков:
>>> 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]])
Как предлагает Тоан, простой способ взломать - сначала выбрать строки, а затем выбрать столбцы поверх них.
>>> 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]])
[Редактировать] Встроенный метод: np.ix_
Недавно я обнаружил, что numpy дает вам встроенный однострочный текст для выполнения именно того , что предлагает @Jaime, но без необходимости использовать синтаксис широковещания (который страдает от недостатка читабельности). Из документов:
Используя ix_, можно быстро создать индексные массивы, которые будут индексировать перекрестный продукт.
a[np.ix_([1,3],[2,5])]возвращает массив[[a[1,2] a[1,5]], [a[3,2] a[3,5]]],
Итак, вы используете это так:
>>> a = np.arange(20).reshape((5,4))
>>> a[np.ix_([0,1,3], [0,2])]
array([[ 0, 2],
[ 4, 6],
[12, 14]])
И способ, которым это работает, состоит в том, что он заботится о выравнивании массивов, как предложил Хайме, чтобы вещание происходило правильно:
>>> np.ix_([0,1,3], [0,2])
(array([[0],
[1],
[3]]), array([[0, 2]]))
Кроме того, как говорит MikeC в комментарии, np.ix_ имеет преимущество в том, что возвращает представление, чего не было в моем первом (до редактирования) ответе. Это означает, что теперь вы можете назначить индексированный массив:
>>> 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]])
ОБЛАСТЬ ПРИМЕНЕНИЯ:
>>> a[[0,1,3]][:,[0,2]]
array([[ 0, 2],
[ 4, 6],
[12, 14]])
ИЛИ ЖЕ:
>>> a[[0,1,3],::2]
array([[ 0, 2],
[ 4, 6],
[12, 14]])
С помощью np.ix_ - самый удобный способ сделать это (как ответили другие), но вот еще один интересный способ сделать это:
>>> rows = [0, 1, 3]
>>> cols = [0, 2]
>>> a[rows].T[cols].T
array([[ 0, 2],
[ 4, 6],
[12, 14]])
np.ix_ делает вычитание матрицы намного проще! Однако мне все еще было любопытно узнать о принятом ответе и хочу знать, почему этот синтаксис работает. Подумав об этом, думаю, у меня есть подсказки. Также здесь я отвечаю на вопрос @Aetos. Чтобы иметь лучшую иллюстрацию, я использую изображение здесь, введите описание изображения здесь. Надеюсь, это поможет.