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

我有一个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。所以我希望输出如下： 3223789342343434

---

在Python中，最简单的方法是什么？假设您将它放在如下的numpy数组中：（忽略科学符号）

你可以做：

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为1到1000，c2为100到1，c3为随机）的数组中，查找每个c1的前5个值只需在我的计算机上花费约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

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