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

Answer 2

为了性能，我建议使用 NumPy -

def mask_variable_largest_per_row(df, n_max):
    a = df.values
    m,n = a.shape

    nan_row_count = np.isnan(a).sum(1)
    n_reset = n-n_max.values-nan_row_count
    n_reset.clip(min=0, max=n-1, out = n_reset)

    sidx = a.argsort(1)
    mask = n_reset[:,None] > np.arange(n)

    c = sidx[mask]
    r = np.repeat(np.arange(m), n_reset)
    a[r,c] = np.nan
    return df

示例运行 -

In [182]: df
Out[182]: 
         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   5.0  24.0  25.0

In [183]: n_max = pd.Series([2,3,2])

In [184]: mask_variable_largest_per_row(df, n_max)
Out[184]: 
       a     b     c     d     e
row1 NaN   NaN   NaN   4.0   5.0
row2 NaN  12.0  13.0  14.0   NaN
row3 NaN   NaN   NaN  24.0  25.0

进一步提升：引入numpy.argpartition 来替换numpy.argsort 应该会有所帮助，因为我们不关心要重置为NaNs 的索引的顺序。因此，基于numpy.argpartition 的将是 -

def mask_variable_largest_per_row_v2(df, n_max):
    a = df.values
    m,n = a.shape

    nan_row_count = np.isnan(a).sum(1)
    n_reset = n-n_max.values-nan_row_count
    n_reset.clip(min=0, max=n-1, out = n_reset)

    N = (n-n_max.values).max()
    N = np.clip(N, a_min=0, a_max=n-1)

    sidx = a.argpartition(N, axis=1) #sidx = a.argsort(1)
    mask = n_reset[:,None] > np.arange(n)

    c = sidx[mask]
    r = np.repeat(np.arange(m), n_reset)
    a[r,c] = np.nan
    return df

运行时测试

其他方法 -

def pandas_rank_based(df, n_max):
    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)
    return df[selected]

验证和时间安排 -

In [387]: arr = np.random.rand(1000,1000)
     ...: arr.ravel()[np.random.choice(arr.size, 10000, replace=0)] = np.nan
     ...: df1 = pd.DataFrame(arr)
     ...: df2 = df1.copy()
     ...: df3 = df1.copy()
     ...: n_max = pd.Series(np.random.randint(0,1000,(1000)))
     ...: 
     ...: out1 = pandas_rank_based(df1, n_max)
     ...: out2 = mask_variable_largest_per_row(df2, n_max)
     ...: out3 = mask_variable_largest_per_row_v2(df3, n_max)
     ...: print np.nansum(out1-out2)==0 # Verify
     ...: print np.nansum(out1-out3)==0 # Verify
     ...: 
True
True

In [388]: arr = np.random.rand(1000,1000)
     ...: arr.ravel()[np.random.choice(arr.size, 10000, replace=0)] = np.nan
     ...: df1 = pd.DataFrame(arr)
     ...: df2 = df1.copy()
     ...: df3 = df1.copy()
     ...: n_max = pd.Series(np.random.randint(0,1000,(1000)))
     ...: 

In [389]: %timeit pandas_rank_based(df1, n_max)
1 loops, best of 3: 559 ms per loop

In [390]: %timeit mask_variable_largest_per_row(df2, n_max)
10 loops, best of 3: 34.1 ms per loop

In [391]: %timeit mask_variable_largest_per_row_v2(df3, n_max)
100 loops, best of 3: 5.92 ms per loop

50x+ 比内置的 pandas 有相当不错的加速！