распространять вычисления вдоль pandas строк данных

Question

Мне нужно распространить вычисление (например, задержку) вдоль pandas строк данных.

Я нашел решение, которое использует метод .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

любое предложение о том, чтобы сделать этот код векторизованным и избежать для l oop?

Answer 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