Давайте посмотрим на три решения и приведем сравнение производительности в конце.
Один из подходов, который пытается остаться рядом с пандами, будет следующим:
def f1(df):
# Group together the elements of df.Sum that might have to be added
pos_groups = (df.Sum <= 0).cumsum()
pos_groups[df.Sum <= 0] = -1
# Create the new column and populate it with what is in df.Sum
df['new_sum'] = df.Sum
# Find the indices of the new column that need to be calculated as a sum
indices = df[df.Quantity > 0].index
for i in indices:
# Find the relevant group of positive integers to be summed, ensuring
# that we only consider those that come /before/ the one to be calculated
group = pos_groups[:i+1] == pos_groups[i]
# Zero out all the elements that will be part of the sum
df.new_sum[:i+1][group] = 0
# Calculate the actual sum and store that
df.new_sum[i] = df.Sum[:i+1][group].sum()
f1(df)
Одно место, где возможноПространство для улучшения может быть в pos_groups[:i+1] == pos_groups[i], который проверяет все i+1 элементы, когда, в зависимости от того, как выглядят ваши данные, возможно, сойдет с рук проверка части из них.Однако, скорее всего, это все еще более эффективно на практике.Если нет, вы можете захотеть выполнить итерацию вручную, чтобы найти группы:
def f2(sums, quantities):
new_sums = np.copy(sums)
indices = np.where(quantities > 0)[0]
for i in indices:
a = i
while sums[a] > 0:
s = new_sums[a]
new_sums[a] = 0
new_sums[i] += s
a -= 1
return new_sums
df['new_sum'] = f2(df.Sum.values, df.Quantity.values)
Наконец, еще раз, в зависимости от того, как выглядят ваши данные, есть неплохой шанс, что последний подход можно улучшить с помощью Numba :
from numba import jit
f3 = jit(f2)
df['new_sum'] = f3(df.Sum.values, df.Quantity.values)
Для данных, представленных в вопросе (которые могут быть слишком малы, чтобы обеспечить правильную картину), тесты производительности выглядят следующим образом:
In [13]: %timeit f1(df)
5.32 ms ± 77.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [14]: %timeit df['new_sum'] = f2(df.Sum.values, df.Quantity.values)
190 µs ± 5.23 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each
In [18]: %timeit df['new_sum'] = f3(df.Sum.values, df.Quantity.values)
178 µs ± 10.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Здесь большая часть времени тратится на обновление фрейма данных.Если бы данные были в 1000 раз больше, решение Numba оказалось бы явным победителем:
In [28]: df_large = pd.concat([df]*1000).reset_index()
In [29]: %timeit f1(df_large)
5.82 s ± 63.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [30]: %timeit df_large['new_sum'] = f2(df_large.Sum.values, df_large.Quantity.values)
6.27 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [31]: %timeit df_large['new_sum'] = f3(df_large.Sum.values, df_large.Quantity.values)
215 µs ± 5.76 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)