如何基于top_n列获取DataFrame每行的前N个值？

原始DataFrame #

import pandas as pd

df1 = pd.DataFrame(data={'1': ['a', 'd', 'g', 'j'], 
                         '2': ['b', 'e', 'h', 'k'], 
                         '3': ['c', 'f', 'i', 'l'],
                         'top_n': [1, 3, 2, 1]},
                   index=pd.Series(['ind1', 'ind2', 'ind3', 'ind4'], name='index'))

输出：

1  2  3  top_n
index
ind1   a  b  c      1
ind2   d  e  f      3
ind3   g  h  i      2
ind4   j  k  l      1

需求 #

根据每行top_n列的数值N，仅保留该行前N个列（列1、2、3）的值，其余列设为NaN，期望结果如下：

1    2    3  top_n
index
ind1   a  NaN  NaN      1
ind2   d    e    f      3
ind3   g    h  NaN      2
ind4   j  NaN  NaN      1

解决方案 #

通过生成布尔掩码的方式实现，代码如下：

# 提取需要处理的列（排除top_n）
target_cols = df1.columns[:-1]
# 为目标列生成1-based的位置序号
col_positions = pd.Series(range(1, len(target_cols)+1), index=target_cols)
# 广播比较，生成掩码：列位置 <= 对应行的top_n值则保留
mask = col_positions <= df1['top_n'].values[:, None]
# 将不满足条件的位置设为NaN
df1[target_cols] = df1[target_cols].where(mask)

# 查看结果
print(df1)

代码逻辑说明 #

1. 先分离出需要处理的列（即除top_n外的列）
2. 给这些列按顺序分配1、2、3的位置序号，和top_n的数值规则匹配（1-based）
3. 利用numpy广播特性，让每行的top_n值与所有列的位置序号比较，得到每行需要保留的列的布尔掩码
4. 使用DataFrame.where()方法，将掩码为False的位置替换为NaN

内容的提问来源于stack exchange，提问作者bltSandwich21