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

我稍微改了一下你的例子，打乱了年龄的顺序：

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

如果你想要的所有年龄都在数据的年龄列表中，你可以先按照年龄这个字段对数据进行排序，然后再用searchsorted()快速找到所有的索引：

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

或者你可以使用np.in1d来获取一个布尔数组，然后用它作为索引：

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

Answer 2

这是一个合理的初步方法：

>>> 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的函数，所以可能会比较慢。我们可以很容易地把它转换成纯numpy的写法，使用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）。

Answer 3

你想创建一个子数组，这个子数组只包含在 desired_ages 中指定的索引对应的值。

在Python中，没有直接的语法可以做到这一点，但使用列表推导式可以很好地完成这个任务：

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

不过，这种方法会导致Python在处理 data.values 中的每个元素时，都要遍历一次 desired_ages，这样效率比较低。如果你能在前面的那一行插入

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]