沿熊猫数据框行传播计算

时间:2020-02-26 17:35:54

标签: python pandas vectorization accumulator propagation

我需要沿着熊猫数据帧行传播计算(例如延迟)。

我找到了一个使用.iterrows()方法并且非常慢的解决方案,因此我想知道是否存在针对此问题的矢量化解决方案,因为我的数据量很大。

这是我的方法:

import pandas as pd
import numpy as np
df = pd.DataFrame(index = ['task_1', 'task_2', 'task_3', 'task_4', 'task_5'], columns=['start_time', 'end_time'], data = [[1,2], [3,4], [6,7], [7,8], [10,11] ] )

# set start delay on task 2
start_delay_on_task_2 = 3
df.loc['task_2', 'start_delay'] = start_delay_on_task_2
df['start_delay'].fillna(0, inplace=True)

# compute buffer between tasks
df['buffer_to_next_task'] = df['start_time'].shift(-1) - df['end_time']

这是df的内容:

        start_time  end_time
task_1  1           2
task_2  3           4
task_3  6           7
task_4  7           8
task_5  10          11

现在是有史以来最糟糕的代码来计算总延迟

df['overall_start_delay'] = df['start_delay']
overall_start_delay_idx = df.columns.get_loc('overall_start_delay')
start_delay_idx = df.columns.get_loc('start_delay')
buffer_to_next_task_idx = df.columns.get_loc('buffer_to_next_task')
for i in range(len(df)):
    overall_delay = None
    if list(df.index)[i] <= 'task_2':
        overall_delay = df.iloc[i, start_delay_idx]
    else:
        overall_delay = max(0, df.iloc[i-1, overall_start_delay_idx] - df.iloc[i-1, buffer_to_next_task_idx])
    df.iloc[i, overall_start_delay_idx] = overall_delay

这里是想要的结果

         start_time end_time start_delay    buffer_to_next_task overall_start_delay
task_1   1          2        0.0            1.0                 0.0
task_2   3          4        3.0            2.0                 3.0
task_3   6          7        0.0            0.0                 1.0
task_4   7          8        0.0            2.0                 1.0
task_5   10         11       0.0            NaN                 0.0

关于将此代码向量化并避免for循环的任何建议?

1 个答案:

答案 0 :(得分:1)

这是一个延迟的解决方案:

total_delays = df.start_delay.cumsum()
(total_delays
 .sub(df.buffer_to_next_task
      .where(total_delays.gt(0),0)
      .cumsum().shift(fill_value=0)
     )
   .clip(lower=0)
)

输出:

task_1    0.0
task_2    3.0
task_3    1.0
task_4    1.0
task_5    0.0
dtype: float64