Блог

Я хотел бы векторизовать этот фрагмент кода python с помощью цикла for, зависящего от текущего состояния для скорости и эффективности.

значения для 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)
  

Это кажется излишним. Ваша state переменная в основном является предыдущим значением df_A['a']*10 . Поэтому мы можем просто использовать shift :

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

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

Вы можете создать класс, в котором state есть переменная класса. Это позволит вам написать функцию, которая может быть передана 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 секунды.