Pandas数据框查找每行特定数量的最大元素

Question

I have a DataFrame:

>>> 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）。请注意第三行的打破平局：

>>> 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()，但在那里无法进行处理并列情况...

- Zhang18

---

“row3” 的第二个元素被设置为 NaN，而第一个元素与其相等。那么选择第二个元素的原因是什么？如果选择其中任何一个元素，会有影响吗？ - Divakar

没关系。这正是我评论.nlargest()行为的原因，它具有随机的平局打破功能。 - Zhang18

.rank有一个method参数，可以处理所有类型的“ties”。 - Scott Boston

@ScottBoston，谢谢。仅凭那个提示就解决了我的问题。我将使用 method='first'。 - Zhang18

由于我们正在寻求性能，你是否查看了这篇帖子 - https://dev59.com/saHia4cB1Zd3GeqPaeOa#44074289? - Divakar

Answer 1

根据@ScottBoston在原帖中的评论，可以使用以下基于排名的掩码来解决此问题：

>>> 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

- Zhang18

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

相比于pandas内置的速度，有着50倍以上的很不错的加速！

- Divakar