Индексирование массива NumPy

Question

Простой вопрос по поводу индексации массива для получения подмножества его значений. Скажем, у меня есть recarray, который содержит возрасты в одном пространстве и соответствующие значения в другом. У меня также есть массив, который является моим желаемым подмножеством возрастов. Вот что я имею в виду:

ages = np.arange(100)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([1,4, 16, 29, 80])

Я пытаюсь сделать что-то вроде этого:

data.values[data.ages==desired_ages]

Но это не работает.

Answer 1

Вы хотите создать подмассив, содержащий только те значения, индексы которых находятся в desired_ages.

В Python нет синтаксиса, прямо соответствующего этому, но списочные выражения могут выполнять довольно хорошую работу:

result = [value for index, value in enumerate(data.values) if index in desired_ages]

Тем не менее, выполнение этого способа приводит к сканированию Python через desired_ages для каждого элемента в data.values, что является медленным.Если бы вы могли вставить

desired_ages = set(desired_ages)

в строку ранее, это улучшило бы производительность.(Вы можете определить, является ли значение в постоянным временем, независимо от размера набора.)

---

Полный пример

import numpy as np

ages = np.arange(100)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([1,4, 16, 29, 80])

result = [value for index, value in enumerate(data.values) if index in desired_ages]
print result

Вывод

[0.45852624094611272, 0.0099713014816563694, 0.26695859251958864, 0.10143425810157047, 0.93647796171383935]

Answer 2

Я немного изменил ваш пример, перетасовал порядок возрастов:

import numpy as np
np.random.seed(0)
ages = np.arange(3,103)
np.random.shuffle(ages)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([4, 16, 29, 80])

Если все элементы требуемые_агенты находятся в data.ages, вы можете сначала отсортировать данные по полю возраста, а затем использовать searchsorted() чтобы быстро найти весь индекс:

data.sort(order="ages") # sort by ages
print data.values[np.searchsorted(data.ages, desired_ages)]

или вы можете использовать np.in1d ​​для массива get bool и использовать его в качестве индекса:

print data.values[np.in1d(data.ages, desired_ages)]

Answer 3

Это разумный первый подход:

>>> bool_indices = reduce(numpy.logical_or, 
                          (data.ages == x for x in desired_ages))
>>> data.values[bool_indices]
array([ 0.63143784,  0.93852927,  0.0026815 ,  0.66263594,  0.2603184 ])

Но здесь используются функции Python, так что, вероятно, медленнее. Мы можем довольно легко перевести его в чистый тупик, используя ix_, чтобы массивы хорошо транслировались друг против друга. (meshgrid с замененными аргументами тоже будет работать, но потребует больше памяти.):

>>> bools_2d = numpy.equal(*numpy.ix_(desired_ages, data.ages))
>>> bool_indices = numpy.logical_or.reduce(bools_2d)
>>> data.ages[bool_indices]
array([ 1,  4, 16, 29, 80])
>>> data.values[bool_indices]
array([ 0.32324063,  0.65453647,  0.9300062 ,  0.34534668,  0.12151951])

См. Также HYRY ответ для потенциально более быстрого решения (с использованием searchsorted) и потенциально более читаемого решения (с использованием in1d).