基于另一列值修改DataFrame指定列的非零值 | Pandas

需求回顾 #

给定如下Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'Names': ['A', 'B', 'C', 'D', 'E'],
    'Value': [100, np.nan, 20, 1, 300],
    'COLA': [0, 0, 3, 0, 3],
    'COLB': [4, 2, 0, 1, 0],
    'COLC': [1, 1, 0, 0, 0]
})

需要对COLA、COLB、COLC列的非零值按以下规则替换：

• 当Value列值>30时，替换为1
• 当Value列值<=30或为NaN时，替换为2

实现代码 #

# 指定需要处理的目标列
target_cols = ['COLA', 'COLB', 'COLC']

# 生成替换规则数组：Value>30则取1，其余情况（含NaN）取2
replace_values = np.where(df['Value'] > 30, 1, 2)

# 对目标列执行替换：非零值用规则值替换，零值保留原数据
df[target_cols] = df[target_cols].mask(df[target_cols] != 0, replace_values[:, None])

# 查看结果
print(df)

代码说明 #

1. 目标列指定：通过target_cols明确要修改的列，避免误操作其他列
2. 替换规则生成：利用np.where生成与Value列对应的替换值，NaN在比较运算中会返回False，自动归为替换为2的场景
3. 批量替换：使用mask方法高效处理批量替换，replace_values[:, None]将一维数组转为二维，实现与目标列的形状匹配（广播机制），仅对非零单元格执行替换

运行结果 #

Names  Value  COLA  COLB  COLC
0     A  100.0     0     1     1
1     B    NaN     0     2     2
2     C   20.0     2     0     0
3     D    1.0     0     2     0
4     E  300.0     1     0     0

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