pandas:仅保留前n个值并将其他值设置为0

Question

在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

Answer 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