Python/ Pandas : Finding index for the nlargest and keeping only those above a value

Question

我想检索与包含 n 个最大值的列集相对应的 n 个列名。然后，仅当值高于给定阈值时，这些名称才会保留在列表中。

举个例子，给定三列“Paul”、“Eric”、“Marc”，假设我最多需要 2 个列名，阈值为 0.5。然后我将获得“最佳”列，如下所示:

import pandas as pd
import numpy as np

start = '2020-01-01 00:00+00:00'
end = '2020-01-01 05:00+00:00'

pr1h = pd.period_range(start=start, end=end, freq='1h')
r = len(pr1h)

df = pd.DataFrame(np.random.rand(r,3), index=pr1h, columns=['Paul', 'Marc', 'Eric'])

处理后:

df
                      Paul      Marc      Eric            Bests
2020-01-01 00:00  0.124974  0.525182  0.415339         ['Marc']
2020-01-01 01:00  0.991917  0.489479  0.668359 ['Paul', 'Eric']
2020-01-01 02:00  0.204156  0.610034  0.644715 ['Eric', 'Marc']
2020-01-01 03:00  0.385546  0.981641  0.089667         ['Marc']
2020-01-01 04:00  0.912330  0.711822  0.148064 ['Paul', 'Marc']
2020-01-01 05:00  0.301186  0.313572  0.323487               []

我可以在 this question/answer 上找到这显示了一种根据给定行中值的排名获取索引的方法。我想这可能是一个起点(可能在速度方面没有优化，因为运行了几次，但这似乎是一个好的开始。

然后我可以:

df1['1st_largest'] = df.columns[df.values.argsort(1)[:,-1]]
df2['2nd_largest'] = df.columns[df.values.argsort(1)[:,-2]]

我的数组不应超过 20 到 50 列，因此我保留使用 argsort 而不是 argpartition。

但现在，我陷入了困境。我不知道如何检查与这些列之一相关的值是否高于 0.5，以便我可以将其放入列表中。

欢迎任何帮助，谢谢!

Answer 1

一种方法是使用 where 和 stack 屏蔽数据帧:

df['Bests'] = (df.where(df.gt(0.5))         # change 0.5 to your threshold
                 .stack().groupby(level=0)
                 .apply(lambda x: x.nlargest(2).index
                                   .get_level_values(1).to_list()
                       )
              )

输出:

                      Paul      Marc      Eric         Bests
2020-01-01 00:00  0.124974  0.525182  0.415339        [Marc]
2020-01-01 01:00  0.991917  0.489479  0.668359  [Paul, Eric]
2020-01-01 02:00  0.204156  0.610034  0.644715  [Eric, Marc]
2020-01-01 03:00  0.385546  0.981641  0.089667        [Marc]
2020-01-01 04:00  0.912330  0.711822  0.148064  [Paul, Marc]
2020-01-01 05:00  0.301186  0.313572  0.323487           NaN