Python 数据帧矢量化 for 循环

Question

我想用 for 循环对这段 python 代码进行矢量化处理，以提高速度和效率。

df_B 的值是根据当前状态 (state) 和对应的 df_A 值计算的。

如有任何想法，我们将不胜感激。

import pandas as pd
df_A = pd.DataFrame({'a': [0, 1, -1, -1, 1, -1, 0, 0] ,})
df_B = pd.DataFrame( data=0, index=df_A.index, columns=['b'])
print(df_A)

state = 0
for index, iter in df_A.iterrows():
    if df_A.loc[index ,'a'] == -1:
        df_B.loc[index ,'b'] = -10 -state
    elif df_A.loc[index, 'a'] == 1:
        df_B.loc[index, 'b'] = 10 - state
    elif df_A.loc[index, 'a'] == 0:
        df_B.loc[index, 'b'] = 0 - state
    temp_state = state
    state += df_B.loc[index, 'b']
print(df_B)

Answer 1

您可以创建一个 class，其中 state 是一个 class 变量。这将允许您编写一个可以提供给 apply 语句的函数。这不是矢量化解决方案，但它比 iterrows 更快。例如：

class ComputeB:
    def __init__(self, state=0):
        self.state = state
    
    def compute_b(self, row):
        row["b"] = row["a"]*10 - self.state
        self.state +=  row["b"]
        return row
df = pd.concat([df_A, df_B], axis = 1)
cb = ComputeB()
df = df.apply(lambda row: cb.compute_b(row), axis = 1)

现在 df["b"] 包含您要计算的值。这确实假定 df_A["a"] 只能包含 0、1 和 -1。在我的机器上有一列 40000 个值，问题中的方法用了 10.4 秒，而这个方法用了 2.95 秒。

Answer 2

这似乎有些矫枉过正。您的 state 变量基本上是 df_A['a']*10 中的先前值。所以我们可以只使用 shift:

s = df_A['a'].mul(10) 

df_B['b'] = s - s.shift(fill_value=0)

Python 数据帧矢量化 for 循环

Python dataframe vectorizing for loop

python

vectorization

pandas