如何在Pandas中按特定条件高效选择并组合列生成金额明细？

原始DataFrame #

import pandas as pd

df = pd.DataFrame(data={
    "id": ['a', 'a', 'b', 'b', 'a', 'c', 'c', 'b'],
    "transaction_amount": [110, 0, 10, 30, 40.4, 62.2, 20, 20],
    "principal_amount":   [100, 0, 0,  0,  40,   60,   0,  0],
    "interest_amount":    [10,  0, 10, 0,  0.4,  0.6,  10, 0],
    "overpayment_amount": [0,   0, 0,  0,  0,    1.6,  10, 20],
})

需求说明 #

生成amount列和transaction_type列，规则如下：

• 若principal_amount、interest_amount、overpayment_amount某列值不为0，为该值单独生成一行，transaction_type对应为principal、interest、overpayment
• 若上述三列值均为0，则取transaction_amount的值作为amount，transaction_type为NaN

期望输出 #

amount transaction_type id
3    30.0              NaN  b
0   100.0        principal  a
4    40.0        principal  a
5    60.0        principal  c
0    10.0         interest  a
2    10.0         interest  b
4     0.4         interest  a
5     0.6         interest  c
6    10.0         interest  c
5     1.6      overpayment  c
6    10.0      overpayment  c
7    20.0      overpayment  b

简洁高效实现方案 #

利用Pandas内置的melt（或stack）进行数据重塑，结合布尔过滤实现需求，避免逐行循环，效率更高：

方法1：使用melt #

# 1. 重塑三列金额数据，过滤非0值
melted_data = df.melt(
    id_vars='id',
    value_vars=['principal_amount', 'interest_amount', 'overpayment_amount'],
    var_name='transaction_type',
    value_name='amount'
).query('amount != 0')

# 2. 简化transaction_type名称
melted_data['transaction_type'] = melted_data['transaction_type'].str.replace('_amount', '')

# 3. 筛选三列金额全为0的行，提取transaction_amount
total_zero_mask = df[['principal_amount', 'interest_amount', 'overpayment_amount']].sum(axis=1) == 0
zero_amount_rows = df.loc[total_zero_mask, ['id', 'transaction_amount']].rename(columns={'transaction_amount': 'amount'})

# 4. 合并两部分结果
final_result = pd.concat([zero_amount_rows, melted_data], ignore_index=True)

方法2：使用stack #

# 1. 将三列金额转为行，过滤非0值
stacked_data = df.set_index('id')[['principal_amount', 'interest_amount', 'overpayment_amount']].stack()
stacked_data = stacked_data[stacked_data != 0].reset_index()
stacked_data.columns = ['id', 'transaction_type', 'amount']
stacked_data['transaction_type'] = stacked_data['transaction_type'].str.replace('_amount', '')

# 2. 处理全0行
total_zero_mask = df[['principal_amount', 'interest_amount', 'overpayment_amount']].sum(axis=1) == 0
zero_amount_rows = df.loc[total_zero_mask, ['id', 'transaction_amount']].rename(columns={'transaction_amount': 'amount'})

# 3. 合并结果
final_result = pd.concat([zero_amount_rows, stacked_data], ignore_index=True)

这两种方法都利用了Pandas的向量化操作和数据重塑特性，比逐行遍历的方式效率提升明显，尤其适合处理大规模数据集。

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