python-为排序的numpy矩阵中第1列的每个值选择第2列的前k个元素

Question

我有一个 numpy 数组，它在前两列上按词法排序，如下所示:

c1  c2  c3
2  0.9  3223  
2  0.8  7899  
2  0.7  23211  
2  0.6  3232  
2  0.5  4478  
1  0.9  342    
1  0.8  3434  
1  0.7  24232   
1  0.6  332  
1  0.5  478

我想要每个 c1 的前两行的 c3 值。所以我想要这样的输出: 3223,7899, 342, 3434

用 Python 最简单的方法是什么

Answer 1

假设你把它放在一个像这样的 numpy 数组中:(忽略科学记数法)

In [86]: arr
Out[86]: 
array([[  1.00000000e+00,   9.00000000e-01,   3.22300000e+03],
       [  1.00000000e+00,   8.00000000e-01,   7.89900000e+03],
       [  1.00000000e+00,   7.00000000e-01,   2.32110000e+04],
       [  1.00000000e+00,   6.00000000e-01,   3.23200000e+03],
       [  1.00000000e+00,   5.00000000e-01,   4.47800000e+03],
       [  2.00000000e+00,   9.00000000e-01,   3.42000000e+02],
       [  2.00000000e+00,   8.00000000e-01,   3.43400000e+03],
       [  2.00000000e+00,   7.00000000e-01,   2.42320000e+04],
       [  2.00000000e+00,   6.00000000e-01,   3.32000000e+02],
       [  2.00000000e+00,   5.00000000e-01,   4.78000000e+02]])

你可以这样做:

arr[np.roll(arr[:,0], k) != arr[:,0],2]

例子:

In [87]: arr[np.roll(arr[:,0], 2) != arr[:,0],2]
Out[87]: array([ 3223.,  7899.,   342.,  3434.])

解释:

我们移动(滚动)k 个位置中的 c1 以获得 c1'。 c1 != c1' 所在的行是 c1 的每个不同值的前 k 行(如果 c1 的值不至少有 k 行，则小于 k 行)。我们使用它来索引原始数组并获得我们想要的 c3 值。

它也应该是完全矢量化的，因此非常高效。 在具有 100000 行和 1000 个不同 c1 值的数组中查找每个 c1 的前 5 个值(对于每个 c1，c1 从 1 到 1000，c2 从 100 到 1，c3 随机)在我的计算机上只需要 ~2.4ms:

In [132]: c1 = np.repeat(np.linspace(1,1000, 1000), 100)

In [133]: c2 = np.tile(np.linspace(100, 1, 100), 1000)

In [134]: c3 = np.random.random_integers(1, 10000, size=100000)

In [135]: arr = np.column_stack((c1, c2, c3))

In [136]: arr
Out[136]: 
array([[  1.00000000e+00,   1.00000000e+02,   2.21700000e+03],
       [  1.00000000e+00,   9.90000000e+01,   9.23000000e+03],
       [  1.00000000e+00,   9.80000000e+01,   1.47900000e+03],
       ..., 
       [  1.00000000e+03,   3.00000000e+00,   7.41600000e+03],
       [  1.00000000e+03,   2.00000000e+00,   2.08000000e+03],
       [  1.00000000e+03,   1.00000000e+00,   3.41300000e+03]])

In [137]: %timeit arr[ np.roll(arr[:,0], 5) != arr[:,0], 2]
100 loops, best of 3: 2.36 ms per loop