如何在 Pandas 数据框中找到每行最大的数值。

Question

给定一个具有$m$列（假设$m=10$）的数据框，对于每一行，我想要找到前$n$个最大的列值（假设$n=2$）。在找到每行的这些前$n$个值之后，我希望将剩余的列值，共计$m-n$个，赋值为0。

例如，从第一个表中提到的数值数据框开始，我试图创建一个符合上述筛选条件的第一个表的表示形式。如果超过$n$列具有相同的值，则较小的列索引号会得到优先考虑。

| col_A | col_B | col_C | col_D | col_E |
|-------|-------|-------|-------|-------|
| 0.1   | 0.1   | 0.3   | 0.4   | 0.5   |
| 0.06  | 0.1   | 0.1   | 0.1   | 0.01  |
| 0.24  | 0.24  | 0.24  | 0.24  | 0.24  |
| 0.20  | 0.25  | 0.30  | 0.12  | 0.02  |

| col_A | col_B | col_C | col_D | col_E |
|-------|-------|-------|-------|-------|
| 0     | 0     | 0     | 0.4   | 0.5   |
| 0     | 0.1   | 0.1   | 0     | 0     |
| 0.24  | 0.24  | 0     | 0     | 0     |
| 0     | 0.25  | 0.3   | 0     | 0     |

有更简单的方法实现吗？向量化的格式可以显著减少在大型数据框上处理的时间

谢谢

- RTM

Answer 1

首先的想法是通过Series.nlargest比较每行中的前N个最大值和使用DataFrame.where设置nset值：

N = 2
df = df.where(df.apply(lambda x: x.eq(x.nlargest(N)), axis=1), 0)
print (df)
   col_A  col_B  col_C  col_D  col_E
0   0.00   0.00    0.0    0.4    0.5
1   0.00   0.10    0.1    0.0    0.0
2   0.24   0.24    0.0    0.0    0.0
3   0.00   0.25    0.3    0.0    0.0

为了提高性能，使用了numpy，来自@Divakar的解决方案：

N = 2
#https://dev59.com/orroa4cB1Zd3GeqPpK-w#61518029
idx = np.argsort(-df.to_numpy(), kind='mergesort')[:,:N]
mask = np.zeros(df.shape, dtype=bool)
np.put_along_axis(mask, idx, True, axis=-1)
df = df.where(mask, 0)
print (df)
   col_A  col_B  col_C  col_D  col_E
0   0.00   0.00    0.0    0.4    0.5
1   0.00   0.10    0.1    0.0    0.0
2   0.24   0.24    0.0    0.0    0.0
3   0.00   0.25    0.3    0.0    0.0

- jezrael

---

这种方法似乎很容易理解。有没有更快的方法来做呢？ - RTM

@codecodercoding - 回答已经被编辑了，Divakar的解决方案对我帮助很大。 - jezrael

谢谢@jezrael。我能理解第二个代码块中使用的逻辑。非常感谢您的帮助。 - RTM