行列计算存在依赖需按序执行，能否摆脱Pandas迭代提速？

当前存在带行依赖的计算场景：col3依赖当前行的col1和col2，col1依赖上一行的col3，这类计算无法直接用Pandas向量化操作。针对50万行6列的DataFrame：

• 直接迭代DataFrame耗时2-3分钟
• 转为Numpy数组迭代后耗时约10秒，需进一步优化提速

示例代码 #

d = {'col1': [1, 0, 0, 0], 'col2': [0, 'a', 'a', 0], 'col3': [0, 0, 0, 0]}
df = pd.DataFrame(data=d, index=[0, 1, 2, 3])

for i in range(df.shape[0]):
    if i > 0:
        df.iloc[i, 0] = df.iloc[i-1, 2]
    if df.iloc[i, 1] == 'a':
        df.iloc[i, 2] = df.iloc[i, 0] + 1
    else:
        df.iloc[i, 2] = df.iloc[i, 0]

优化方案 #

1. Numba加速Numpy迭代

Numba能把Python函数编译为机器码，大幅降低循环开销。针对当前逻辑，用Numba装饰器处理数组：

import numba
import numpy as np
import pandas as pd

@numba.jit(nopython=True)  # nopython模式性能最优
def compute(arr):
    rows = arr.shape[0]
    for i in range(rows):
        if i > 0:
            arr[i, 0] = arr[i-1, 2]
        if arr[i, 1] == 'a':
            arr[i, 2] = arr[i, 0] + 1
        else:
            arr[i, 2] = arr[i, 0]
    return arr

# 转换为numpy数组（根据实际数据类型调整，这里用object兼容字符串）
df_arr = df.to_numpy(dtype=object)
df_arr = compute(df_arr)
# 转回DataFrame
result_df = pd.DataFrame(df_arr, columns=df.columns, index=df.index)

优化点：如果col2是固定类别，把'a'映射为整数（比如1），其他值映射为0，避免Numba处理字符串的额外开销，能再提一波速度。

2. 手动向量化（针对固定逻辑推导）

如果业务逻辑可拆解为数学规律，比如观察到：

• 初始col3 = col1[0]
• 当col2为'a'时，col3 = 上一行col3 +1；否则等于上一行col3
• col1就是上一行的col3

可以直接用Pandas的cumsum和shift实现全向量化计算，这是速度最快的方案：

# 设置初始值
df['col3'] = df['col1'].iloc[0]
# 计算每一行的增量：col2为'a'则加1，否则加0
increments = (df['col2'] == 'a').astype(int)
# 累加增量得到col3
df['col3'] = df['col3'].iloc[0] + increments.cumsum()
# col1取上一行的col3，shift填充初始值
df['col1'] = df['col3'].shift(fill_value=df['col1'].iloc[0])

3. Cython编写扩展（极致性能需求）

如果Numba仍达不到性能要求，可以用Cython把循环逻辑写成C级代码，编译后调用，性能接近纯C。但需要额外的编译步骤，适合复杂且对性能极致要求的场景。

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