Pandas 数据框使用特定于行的 N 查找每行的最大 N 个元素

Question

我有一个数据框:

>>> df = pd.DataFrame({'row1' : [1,2,np.nan,4,5], 'row2' : [11,12,13,14,np.nan], 'row3':[22,22,23,24,25]}, index = 'a b c d e'.split()).T
>>> df
         a     b     c     d     e
row1   1.0   2.0   NaN   4.0   5.0
row2  11.0  12.0  13.0  14.0   NaN
row3  22.0  22.0  23.0  24.0  25.0

和一个系列，指定我想要从每行中获得的前 N ​​个值的数量

>>> n_max = pd.Series([2,3,4])

Panda 使用 df 和 n_max 查找每个元素中最大的 N 个元素的方法是什么(通过随机选择打破平局，就像 .nlargest( ) 会吗)？

所需的输出是

         a     b     c     d     e
row1   NaN   NaN   NaN   4.0   5.0
row2   NaN  12.0  13.0  14.0   NaN
row3  22.0   NaN  23.0  24.0  25.0

我知道如何在所有行中使用统一/固定的 N(例如，N=4)来执行此操作。请注意第 3 行中的平局:

>>> df.stack().groupby(level=0).nlargest(4).unstack().reset_index(level=1, drop=True).reindex(columns=df.columns)
         a     b     c     d     e
row1   1.0   2.0   NaN   4.0   5.0
row2  11.0  12.0  13.0  14.0   NaN
row3  22.0   NaN  23.0  24.0  25.0

但目标还是要有特定于行的N。循环遍历每一行显然不算数(出于性能原因)。我尝试过将 .rank() 与掩码一起使用，但打破平局在那里不起作用......

Answer 1

根据 @ScottBoston 对 OP 的评论，可以使用以下基于排名的掩码来解决此问题:

>>> n_max.index = df.index
>>> df_rank = df.stack(dropna=False).groupby(level=0).rank(ascending=False, method='first').unstack()
>>> selected = df_rank.le(n_max, axis=0)
>>> df[selected]
         a     b     c     d     e
row1   NaN   NaN   NaN   4.0   5.0
row2   NaN  12.0  13.0  14.0   NaN
row3  22.0   NaN  23.0  24.0  25.0