如何向量化处理依赖前一行的DataFrame资产负债表类数据计算？

核心思路：推导递推公式，转化为累积和操作 #

你的场景本质是链式依赖的累积计算，但可以通过数学推导打破循环，用纯向量化操作替代for循环：

根据你的逻辑，我们可以展开递推关系：

• 初始行（t=0）：close_balance[0] = open_balance[0] + additions[0] - subtractions[0]
• 对于t≥1：
  1. open_balance[t] = close_balance[t-1]
  2. 代入第二个公式得：close_balance[t] = close_balance[t-1] + additions[t] - subtractions[t]

这意味着close_balance是初始值加上从第1行到当前行的(additions - subtractions)累积和，完全可以用pandas的内置累积函数实现。

完整向量化实现代码 #

import pandas as pd
import numpy as np

# 示例DataFrame（替换为你的实际数据）
df = pd.DataFrame({
    'open_balance': [100, np.nan, np.nan, np.nan],
    'additions': [10, 20, 30, 15],
    'subtractions': [5, 15, 25, 5]
})

# 检查初始行条件：第0行open_balance不为0且非NaN
initial_open = df.loc[0, 'open_balance']
if not (pd.isna(initial_open) or initial_open == 0):
    # 计算初始close_balance
    df.loc[0, 'close_balance'] = initial_open + df.loc[0, 'additions'] - df.loc[0, 'subtractions']
    
    # 计算增量列：additions与subtractions的差值
    delta = df['additions'] - df['subtractions']
    
    # 批量计算后续行的close_balance：初始值 + 增量的累积和
    df.loc[1:, 'close_balance'] = df.loc[0, 'close_balance'] + delta[1:].cumsum()
    
    # 批量计算open_balance：后续行取上一行的close_balance，初始行保留原值
    df['open_balance'] = df['close_balance'].shift(1).fillna(initial_open)
else:
    # 初始行不满足条件时的处理（可根据需求调整）
    print("初始行open_balance为0或NaN，跳过计算")

print(df)

关键优势 #

1. 完全向量化：没有for循环，利用pandas的内置优化函数，计算效率比循环高几个数量级
2. 避免警告：直接对原DataFrame或明确复制的对象操作，消除SettingWithCopyWarning
3. 可扩展性：如果需要按分组计算，只需套上groupby逻辑即可（比如df.groupby('group_col').apply(your_calculation_func)）

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