Question

我有一个类似于：

的数据框

device_id   s2  s41 s47 s14
30          0   0   0   0.003
125         0   0   0   0
32          0   0   0   0
45          0   0   0   0

目标是从每一行获得3个最高s，如果匹配，则选择最大索引。例如，对于第一行，s2，s41和s47之间存在匹配，因此将与s41一起选择s47。因此，第1行的选择将是s14，s47和s41。有超过200万条记录和250多个记录，因此我需要一种有效的方法。我已尝试使用iterrows，然后对每一行进行排序，但速度很慢，所有数据都需要一个多小时。

最终目标是在字典中搜索s（s14，s47，s41），其中这些值是关键字并从中获取适当的值。

有人可以帮助我有效地做到这一点。谢谢

Answer 1

我会使用heapq's nlargest：

In [11]: df
Out[11]:
   device_id  s2  s41  s47    s14
0         30   0    0    0  0.003
1        125   0    0    0  0.000
2         32   0    0    0  0.000
3         45   0    0    0  0.000

In [12]: nlargest(3, df.columns[1:], key=lambda x: int(x[1:]))
Out[12]: ['s47', 's41', 's14']

In [13]: df[["device_id"] + nlargest(3, df.columns[1:], key=lambda x: int(x[1:]))]
Out[13]:
   device_id  s47  s41    s14
0         30    0    0  0.003
1        125    0    0  0.000
2         32    0    0  0.000
3         45    0    0  0.000

注意：如果device_id是索引，那就更容易了：

In [21]: df1
Out[21]:
           s2  s41  s47    s14
device_id
30          0    0    0  0.003
125         0    0    0  0.000
32          0    0    0  0.000
45          0    0    0  0.000

In [22]: df1[nlargest(3, df1.columns, key=lambda x: int(x[1:]))]
Out[22]:
           s47  s41    s14
device_id
30           0    0  0.003
125          0    0  0.000
32           0    0  0.000
45           0    0  0.000

Answer 2

我不知道熊猫，但我知道它是基于numpy的，所以这是一个numpy解决方案。它使用argpartition有效地获取每行中最大4的索引。不幸的是，numpy使用的算法并不稳定，所以如果这四个中最小的两个相等，我们必须做一个完整的排序，排序让我们选择一个稳定的算法。

代码（无法检查我的装备上的2m行b / c内存，但0.5m大约需要2秒左右）：

import numpy as np

def stable_high_3(data):
    n, m = data.shape
    high_4 = np.argpartition(data, np.arange(m-4, m), axis=-1)[:, -4:]
    must_check = np.where(data[np.arange(n), high_4[:, 0]]
                          == data[np.arange(n), high_4[:, 1]])[0]
    high_4[must_check, -3:] = np.argsort(data[must_check], axis=-1,
                                         kind='mergesort')[:, -3:]
    return high_4[:, -3:]

data = np.reshape(np.arange(30)%5, (-1, 6))
print(data)
print(stable_high_3(data))

data = np.reshape(np.arange(256*2**18)%50, (-1, 256))
stable_high_3(data)

打印

[[0 1 2 3 4 0]
 [1 2 3 4 0 1]
 [2 3 4 0 1 2]
 [3 4 0 1 2 3]
 [4 0 1 2 3 4]]
[[2 3 4]
 [1 2 3]
 [5 1 2]
 [0 5 1]
 [4 0 5]]

对每一行进行排序，并从数据框中选择前三个值

2 个答案: