使用索引值从numpy数组中选择项目
假设我有一对numpy数组，X和I，它们的样子是这样的（X是二维的，I是一维的）
X I
-----------------
3.4 9.13 0
3.5 3.43 1
3.6 2.01 2
3.7 6.11 0
3.8 4.95 1
3.9 7.02 2
4.0 4.41 3
4.1 0.23 0
4.2 0.99 1
4.3 1.02 0
4.4 5.61 1
4.5 7.55 2
4.6 8.10 0
4.7 0.33 2
4.8 0.80 1
我想做两件事：
Y = indexby(X,I,I0)：给定一个值I0，找到X中与I中匹配的行。例如，如果I是2，我想找到以下数组：3.6 2.01 3.9 7.02 4.5 7.55 4.7 0.33Y = indexby(X,I)：返回一个字典，里面包含所有可能的键k，使得Y[k] == indexby(X,I,k)。在我的示例数据中，这将产生以下结果：Y[0] = 3.4 9.13 3.7 6.11 4.1 0.23 4.3 1.02 4.6 8.10 Y[1] = 3.5 3.43 3.8 4.95 4.2 0.99 4.4 5.61 4.8 0.80 Y[2] = 3.6 2.01 3.9 7.02 4.5 7.55 4.7 0.33 Y[3] = 4.0 4.41
有没有numpy的函数可以做到这一点？我不太确定该找什么，所以很难找到它们。
我知道我可以手动完成这个，但出于性能考虑，我想使用内置的numpy函数，因为我应用中的数组通常有10万到100万行。
3 个回答
首先，我会给你展示一个使用 结构化数组 的好方法。链接中的文档提供了很多关于如何索引、排序和创建结构化数组的有用信息。
让我们先定义一部分你的数据，
import numpy as np
X = np.array( [[3.4,9.13], [3.5,3.43], [3.6,2.01], [3.7,6.11],
[3.8,4.95], [3.9,7.02], [4.0,4.41]] )
I = np.array( [0,1,2,0,1,2,3], dtype=np.int32 )
结构化数组
如果我们用这些数据创建一个结构化数组（也就是一个结构体数组），那么问题就变得简单了，
sa = np.zeros( len(X), dtype=[('I',np.int64),('X',np.float64,(2))] )
在这里，我们创建了一个空的结构化数组。数组中的每个元素都是一个64位整数和一个包含两个64位浮点数的数组。传给 dtype 的列表定义了这个结构体，每个元组代表结构体的一个部分。元组里包含一个标签、一个类型和一个形状。形状部分是可选的，默认是标量（单个值）。
接下来，我们用你的数据填充这个结构化数组，
sa['I'] = I
sa['X'] = X
此时，你可以这样访问记录，
>>> sa['X'][sa['I']==2]
array([[ 3.6 , 2.01],
[ 3.9 , 7.02]])
在这里，我们请求所有的 'X' 记录，并使用通过 sa['I']==2 语句创建的布尔数组进行索引。然后可以使用列表推导式构建你想要的字典，
d = { i:sa['X'][sa['I']==i] for i in np.unique(sa['I']) }
接下来是两个使用标准 numpy 数组的解决方案。第一个使用 np.where，保持数组不变，另一个则涉及对数组进行排序，这在处理大规模的 I 时应该更快。
使用 np.where
使用 np.where 并不是绝对必要的，因为数组可以使用从 I==I0 产生的布尔数组进行索引，但在某些情况下，拥有实际的整数索引是很有用的。
def indexby1( X,I,I0 ):
indx = np.where( I==I0 )
sub = X[indx[0],:]
return sub
def indexby2( X,I ):
d = {}
I0max = I.max()
for I0 in range(I0max+1):
d[I0] = indexby1( X, I, I0 )
return d
d = indexby2( X, I )
排序和提取数据块
另外，你也可以使用之前提到的排序方案，直接返回数据块，
def order_arrays( X, I ):
indx = I.argsort()
I = I[indx]
X = [indx] # equivalent to X = X[indx,:]
return X, I
def indexby(X, I, I0=None):
if I0 == None:
d = {}
for I0 in range(I.max()+1):
d[I0] = indexby( X, I, I0 )
return d
else:
ii = I.searchsorted(I0)
ff = I.searchsorted(I0+1)
sub = X[ii:ff]
return sub
X,I = order_array( X, I )
d = indexby( X, I )
在这里，我把之前的两个函数合并成一个递归函数，正如你在问题中描述的那样。这当然会修改原始数组。
如果你想试试 pandas，它在处理 groupby 数据方面非常强大。下面是你可以实现你所需功能的方法：
In [34]: import numpy as np
In [35]: import pandas as pd
#I defined you X, I already
In [36]: X
Out[36]:
array([[ 3.4 , 9.13],
[ 3.5 , 3.43],
[ 3.6 , 2.01],
[ 3.7 , 6.11],
[ 3.8 , 4.95],
[ 3.9 , 7.02],
[ 4. , 4.41],
[ 4.1 , 0.23],
[ 4.2 , 0.99],
[ 4.3 , 1.02],
[ 4.4 , 5.61],
[ 4.5 , 7.55],
[ 4.6 , 8.1 ],
[ 4.7 , 0.33],
[ 4.8 , 0.8 ]])
In [37]: I
Out[37]: array([0, 1, 2, 0, 1, 2, 3, 0, 1, 0, 1, 2, 0, 2, 1], dtype=int64)
In [38]: dataframe=pd.DataFrame (data=X, index=I, columns=['X1','X2'])
In [39]: dataframe.index.name='I' #This is not necessary
In [40]: print dataframe
X1 X2
I
0 3.4 9.13
1 3.5 3.43
2 3.6 2.01
0 3.7 6.11
1 3.8 4.95
2 3.9 7.02
3 4.0 4.41
0 4.1 0.23
1 4.2 0.99
0 4.3 1.02
1 4.4 5.61
2 4.5 7.55
0 4.6 8.10
2 4.7 0.33
1 4.8 0.80
这段代码定义了一个数据框，其中 I 是索引，X 是数据。现在，如果你需要 I=2 的行，可以这样做：
In [42]: print dataframe.ix[2]
X1 X2
I
2 3.6 2.01
2 3.9 7.02
2 4.5 7.55
2 4.7 0.33
如果你想列出所有的组：
In [43]: for i, grouped_data in dataframe.groupby(level='I'): #without level=, you can group by a regular column like X1
....: print i
....: print grouped_data
....:
0
X1 X2
I
0 3.4 9.13
0 3.7 6.11
0 4.1 0.23
0 4.3 1.02
0 4.6 8.10
1
X1 X2
I
1 3.5 3.43
1 3.8 4.95
1 4.2 0.99
1 4.4 5.61
1 4.8 0.80
2
X1 X2
I
2 3.6 2.01
2 3.9 7.02
2 4.5 7.55
2 4.7 0.33
3
X1 X2
I
3 4 4.41
如果你只想查看每个组的统计信息，可以这样做：
In [47]: print dataframe.groupby(level='I').sum() #try other funcs like mean, var, .
X1 X2
I
0 20.1 24.59
1 20.7 15.78
2 16.7 16.91
3 4.0 4.41
这里有一些更高级的函数，但我们先来看看怎么用库里最简单的东西来实现，因为你每天都会用到这些简单的函数。
>>> matches = (I == 2)
>>> matches
array([False, False, True, False, False, True, False, False, False,
False, False, True, False, True, False], dtype=bool)
>>> indices = np.nonzero(matches)
>>> indices
(array([ 2, 5, 11, 13]),)
>>> xvals = X[indices]
>>> xvals
array([[ 3.6 , 2.01],
[ 3.9 , 7.02],
[ 4.5 , 7.55],
[ 4.7 , 0.33]])
最后一步可能看起来有点复杂。想了解更多，可以查看索引的教程。
一旦你明白了==这个操作符和nonzero是怎么工作的，可以看看和nonzero在同一部分的其他函数，你会发现有两种更简短的方法来实现这个。