使用索引值从numpy数组中选择项

Question

假设我有一对numpy数组X和I，看起来像这样（X是2D，I是1D）

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

我想做两件事：

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.33  
   

2. Y = 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函数，因为我的应用程序中的数组的行数通常在100000-1000000范围内。你知道吗

Answer 1

如果您想尝试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