如何使用Pandas消除for循环并优化DataFrame处理代码

假设我们有一个形状为(1000,5)的DataFrame df：

step 1   step 2   step 3    step 4   step 5
300       200      500       300      3000  
400       700      1500      2000     2100 
3000      1500     500       5000     4000 
...        ...        ...    ...       ...

还有一个形状为(1000,1)的Series max_cash：

Max Cash
1500
3000
5000
...

我们有一个future_cash()函数用于获取下一步要添加的现金，规则为：若原金额与待加现金之和大于等于max_cash对应值，则设为0；否则添加该现金（示例每次加1000，预期结果如下）：

step 1   step 2   step 3     step 4  step 5
1300      1200       0       1300      0  
1400      1700       2500    0         0 
4000      2500       3500    0         0 
...        ...        ...    ...       ...

当前采用双重循环实现：

# 遍历DataFrame的列
for i in range (0,5):
    deposit = future_cash()
    df.iloc[:, i] = fix_values_for_steps (df.iloc[:, i],max_cash, deposit)

# 自定义函数，遍历行判断并更新值
def fix_values_for_steps(col_from_df: pd.Series, maximum_cash: pd.Series, money:int) -> pd.Series:
    for i in range(len(df)):   
        if df['Maximum Cash'].iloc[i] >=  col_from_df.iloc[i]+money:
              col_from_df.iloc[i] = col_from_df.iloc[i]+money   
        else:
              col_from_df.iloc[i] = 0

请问如何优化这段代码？是否可以仅用单循环处理整列？

完全可以去掉内层循环，利用pandas的向量化操作处理整列，只保留外层的列循环（因为每列需要单独调用future_cash()获取不同的deposit值）。

优化后的代码 #

首先重构fix_values_for_steps函数，用向量化替代行循环：

def fix_values_for_steps(col_from_df: pd.Series, maximum_cash: pd.Series, money: int) -> pd.Series:
    # 计算原数值加存款后的结果
    added = col_from_df + money
    # 条件判断：如果added >= maximum_cash则设为0，否则保留added
    return added.where(added < maximum_cash, 0)

然后外层循环可以简化为遍历列名（更直观）：

for col in df.columns:
    deposit = future_cash()
    df[col] = fix_values_for_steps(df[col], max_cash['Max Cash'], deposit)

优化原理 #

• pandas的Series支持逐元素的算术运算和条件判断，added.where(cond, other)会自动遍历整个Series，满足条件的保留原值，不满足的替换为指定值，完全替代了手动行循环。
• 这种向量化操作比Python级别的循环快得多，尤其是处理1000行以上的数据时，性能提升会非常明显。

进阶优化（若条件允许） #

如果future_cash()可以一次性生成所有5列的deposit值（比如返回一个长度为5的列表），那么连外层列循环也可以去掉，直接用DataFrame级别的向量化操作：

# 假设future_cash()返回[deposit1, deposit2, deposit3, deposit4, deposit5]
deposits = future_cash()
# 将deposits转为与df列数匹配的DataFrame，自动广播到每一行
deposit_df = pd.DataFrame([deposits]*len(df), columns=df.columns)
# 计算所有列的加存款后的值
added = df + deposit_df
# 批量条件替换：added >= max_cash对应值则设为0，否则保留added
df = added.where(added < max_cash['Max Cash'].values[:, None], 0)

但这个方案的前提是future_cash()能一次性生成所有列的deposit，否则外层的列循环是必要的——因为每列需要独立获取deposit值。

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