如何根据Dict列条件调整DataFrame中A、B列的数值位置

问题描述 #

现有一个包含A、B、Dict三列的DataFrame：

• Dict列的值仅为"A"或"B"，用于标记该行数值应归属的目标列
• A和B列存储数值型数据（支持负数）

原数据示例 #

A            B             Dict
0           123            B
0           123            A
123         0              A

需求说明 #

根据Dict列的值调整A、B列的数值：

• 若Dict值为"A"：将该行B列的值移至A列，B列设为0
• 若Dict值为"B"：将该行A列的值移至B列，A列设为0
• 最终要求每行数值仅存在于Dict指示的列，另一列为0

目标数据示例 #

A            B             Dict
0           123             B
123          0              A
123          0              A

实现方法 #

方法1：使用loc索引批量处理 #

这是效率较高的向量式操作，适合大数据集：

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    'A': [0, 0, 123],
    'B': [123, 123, 0],
    'Dict': ['B', 'A', 'A']
})

# 处理Dict为"A"的行
df.loc[df['Dict'] == 'A', 'A'] += df.loc[df['Dict'] == 'A', 'B']
df.loc[df['Dict'] == 'A', 'B'] = 0

# 处理Dict为"B"的行
df.loc[df['Dict'] == 'B', 'B'] += df.loc[df['Dict'] == 'B', 'A']
df.loc[df['Dict'] == 'B', 'A'] = 0

方法2：使用apply逐行处理 #

适合逻辑更复杂的场景，小数据集下更直观：

def adjust_row(row):
    target_col = row['Dict']
    if target_col == 'A':
        row['A'] = row['A'] + row['B']
        row['B'] = 0
    elif target_col == 'B':
        row['B'] = row['A'] + row['B']
        row['A'] = 0
    return row

df = df.apply(adjust_row, axis=1)

方法3：使用numpy.where简化代码 #

用向量式操作实现，代码更简洁：

import numpy as np

# 先复制原列值避免覆盖
temp_a = df['A'].copy()
temp_b = df['B'].copy()

df['A'] = np.where(df['Dict'] == 'A', temp_a + temp_b, 0)
df['B'] = np.where(df['Dict'] == 'B', temp_a + temp_b, 0)

验证结果 #

执行上述任意方法后，DataFrame都会符合目标数据的要求，确保数值仅保留在Dict指定的列中。

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