【问题标题】：pandas: Keep only top n values and set others to 0pandas：只保留前 n 个值并将其他值设置为 0
【发布时间】：2018-11-06 09:49:05
【问题描述】：

在 pandas 数据框中，对于每一行，我只想保留前 N 个值并将其他所有值设置为 0。我可以遍历行并做到这一点，但我确信 python/pandas 可以在一行中优雅地做到这一点。

例如：对于 N = 2

Input:
A   B   C   D
4   10  10  6
5   20  50  90
6   30  6   4
7   40  12  9

Output:
A   B   C   D
0   10  10  0
0   0   50  90
6   30  6   0
0   40  12  0

【问题讨论】：

请解释一下输出。对于列A，您将7 替换为0 并保留6，这是怎么回事？
@MayankPorwal 每个 row 似乎只保留最大的 2 个值（或更多，如果有平局）
是的，没错。将进行编辑以使其清晰。

标签： python pandas

【解决方案1】：

使用rank 和参数axis=1 和method='min' 和ascending=False 作为：

N = 2
df = df.mask(df.rank(axis=1, method='min', ascending=False) > N, 0)

或者使用np.where 和pd.DataFrame 比mask 方法快：

df = pd.DataFrame(np.where(df.rank(axis=1,method='min',ascending=False)>N, 0, df),
                  columns=df.columns)

print(df)
   A   B   C   D
0  0  10  10   0
1  0   0  50  90
2  6  30   6   0
3  0  40  12   0

说明：

第 1 步： 首先，我们需要找到该行中最小的 2 个数字是什么，以及是否需要考虑重复项。因此，使用axis=1 跨行排名和重复值将由method='min' 和ascending = False 处理：

print(df.rank(axis=1, method='min', ascending=False))
     A    B    C    D
0  4.0  1.0  1.0  3.0
1  4.0  3.0  2.0  1.0
2  2.0  1.0  2.0  4.0
3  4.0  1.0  2.0  3.0

第 2 步： 其次，我们需要根据条件过滤值大于 (N) 的位置，然后使用 mask 更改这些值：

print(df.rank(axis=1, method='min', ascending=False) > N)
       A      B      C      D
0   True  False  False   True
1   True   True  False  False
2  False  False  False   True
3   True  False  False   True

print(df.mask(df.rank(axis=1, method='min', ascending=False) > N, 0))
   A   B   C   D
0  0  10  10   0
1  0   0  50  90
2  6  30   6   0
3  0  40  12   0

【讨论】：

@RavinderSingh13 添加解释检查一下。
当我将它应用于我的完整数据集时，这不起作用。您可以检查更多列。
@NikhilUtane 你能显示它失败的样本数据吗？
好的，我复制了一部分样本数据here
是的，这就像一个魅力，几乎是瞬间的，而其他解决方案大约需要 10 秒。请更新原始答案。

【解决方案2】：

用途：

N = 2
df = df.where(df.apply(lambda x: x.isin(x.nlargest(N)), axis=1), 0)
print (df)
   A   B   C   D
0  0  10  10   0
1  0   0  50  90
2  6  30   6   0
3  0  40  12   0

或者：

import heapq
N = 2
df = df.where(df.apply(lambda x: x.isin(heapq.nlargest(N, x)), axis=1), 0)
print (df)
   A   B   C   D
0  0  10  10   0
1  0   0  50  90
2  6  30   6   0
3  0  40  12   0

【讨论】：

输出不太正确。 OP 预期输出中的索引 2 是 6 30 6 0

【解决方案3】：

使用 nlargest 得到 N 个最大的数：

df.mask(~df.apply(lambda x: x.isin(x.nlargest(2)), axis=1), 0)

输出：

    A   B   C   D
0   0   10  10  0
1   0   0   50  90
2   6   30  6   0
3   0   40  12  0

【讨论】：

【解决方案4】：

您可以通过np.apply_along_axis 使用scipy.stats.rankdata，并提供给pd.DataFrame.where：

from scipy.stats import rankdata

df[:] = df.where(np.apply_along_axis(rankdata, 1, df, method='max') > 2, 0)

print(df)

   A   B   C   D
0  0  10  10   0
1  0   0  50  90
2  6  30   6   0
3  0  40  12   0

性能基准测试

pd.DataFrame.rank 是以下最有效的解决方案； apply + lambda 表现最差。

from scipy.stats import rankdata
from heapq import nlargest

df = pd.concat([df]*100, ignore_index=True)

%timeit df.mask(df.rank(axis=1, method='min', ascending=False) > 2, 0)       # 2.23 ms per loop
%timeit df.where(np.apply_along_axis(rankdata, 1, df, method='max') > 2, 0)  # 45 ms per loop
%timeit df.where(df.apply(lambda x: x.isin(nlargest(2, x)), axis=1), 0)      # 92.4 ms per loop
%timeit df.mask(~df.apply(lambda x: x.isin(x.nlargest(2)), axis=1), 0)       # 274 ms per loop

【讨论】：