Fancy-индексация — NumPy

Коротко

a[[0, 2, 4]] — элементы с этими индексами. Порядок любой, повторы и отрицательные индексы разрешены.
Результат — копия (в отличие от среза), форма результата равна форме массива индексов.
В двумерном массиве a[[i1, i2], [j1, j2]] берёт пары (i1, j1), (i2, j2), а не подматрицу. Подматрица — a[np.ix_(строки, столбцы)].
Типичный источник индексов — argsort, argmax, nonzero.
a[idx] += 1 при повторах в idx прибавит один раз; для накопления — np.add.at.

Синтаксис

a[[0, 2, 4]]                   # список или массив индексов
a[idx] = 0                     # запись по индексам
m[[0, 2]]                      # строки 0 и 2
m[[0, 2], [1, 3]]              # элементы (0, 1) и (2, 3)
m[np.ix_([0, 2], [1, 3])]      # подматрица 2 × 2
np.add.at(a, idx, 1)           # прибавить с учётом повторов

Примеры

Данные для примеровиспользуются во всех примерах ниже

prices = np.array([250, 120, 480, 90, 310])   # цены пяти товаров
table = np.array([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

Несколько элементов по индексам

prices[[0, 2, 4]]

Вывод

array([250, 480, 310])

Порядок, повторы и argsort

print(prices[[4, 0, 0, -1]])        # любой порядок, повторы, отрицательные
order = np.argsort(prices)          # индексы по возрастанию цены
prices[order]

Вывод

[310 250 250 310]
array([ 90, 120, 250, 310, 480])

Строки матрицы

table[[2, 0]]                       # строки 2 и 0

Вывод

array([[7, 8, 9],
       [1, 2, 3]])

Пары индексов и подматрица

print(table[[0, 2], [1, 2]])        # пары: (0, 1) и (2, 2)
table[np.ix_([0, 2], [1, 2])]       # подматрица: строки 0, 2 × столбцы 1, 2

Вывод

[2 9]
array([[2, 3],
       [8, 9]])

Запись по индексам

discounted = prices.copy()
discounted[[1, 3]] = 0
discounted

Вывод

array([250,   0, 480,   0, 310])

Подводные камни

Результат — копия, а не представление

Изменения выбранного массива не попадают в исходный. Чтобы изменить исходный, записывайте прямо в него: a[idx] = ....

Изменение выбранных элементов

chosen = prices[[0, 1]]
chosen[0] = -1
prices                              # исходный массив не изменился

Вывод

array([250, 120, 480,  90, 310])

Повторяющиеся индексы при +=

a[idx] += 1 выполняется как a[idx] = a[idx] + 1: при повторах одно и то же значение записывается несколько раз, и прибавка засчитывается однажды.

Счётчики по индексам

counts = np.zeros(3, dtype=int)
counts[[0, 0, 2]] += 1              # индекс 0 дважды, но прибавится один раз
print(counts)
counts = np.zeros(3, dtype=int)
np.add.at(counts, [0, 0, 2], 1)     # каждое вхождение учитывается
counts

Вывод

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