условное сложение в пандах

Question

У меня есть следующий фрейм данных в пандах

   Code      Sum      Quantity
   0         -12      0
   1          23      0
   2         -10      0
   3         -12      0
   4         100      0
   5         102      201
   6          34      0
   7         -34      0
   8         -23      0
   9         100      0
   10        100      0
   11        102      300

Мой желаемый фрейм данных будет

  Code      Sum      Quantity    new_sum
   0         -12      0          -12
   1          23      0           23
   2         -10      0          -10
   3         -12      0          -12
   4         100      0           0
   5         102      201         202 
   6          34      0           34
   7         -34      0          -34
   8         -23      0          -23
   9         100      0           0
   10        100      0           0
   11        102      300         302

Логика

Сначала я проверю ненулевые значения в количествеВ приведенных выше примерах данных мы получили первое ненулевое вхождение числа в index 4 which is 201, затем я хочу добавить столбец суммы, пока я не получу положительные значения в строке выше index 4

Я написал код, который используетцикл, но тогда мне нужно сканировать более 1 миллиона строк, и это не дает мне желаемый результат.

for i in range(len(final_df)):
   if(final_df['Quantity'][i] != 0):
      final_df['new_sum'][i] = final_df['Sum'][i].shift(1).sum()
   else:
      final_df['new_sum'][i] = final_df['Sum'][i] 

Answer 1

Отредактированный ответ после уточнения ...

Этот будет немного медленнее из-за понимания списка и цикла.

Настройка:

import pandas as pd
import numpy as np

data = [[  0, -12,   0],
        [  1,  23,   0],
        [  2, -10,   0],
        [  3, -12,   0],
        [  4, 100,   0],
        [  5, 102, 201],
        [  6,  34,   0],
        [  7, -34,   0],
        [  8, -23,   0],
        [  9, 100,   0],
        [ 10, 100,   0],
        [ 11, 102, 300]]

df = pd.DataFrame(data, columns=['Code', 'Sum', 'Quantity'])

print(df)

    Code  Sum  Quantity
0      0  -12         0
1      1   23         0
2      2  -10         0
3      3  -12         0
4      4  100         0
5      5  102       201
6      6   34         0
7      7  -34         0
8      8  -23         0
9      9  100         0
10    10  100         0
11    11  102       300

Код:

# copy columns from input dataframe and invert
df1 = df[['Sum', 'Quantity']][::-1].copy()

# make an array to hold result column values
new_sum_array = np.zeros(len(df1)).astype(int)
df_sum = df1.Sum.values

# locate the indices of the pos values in "Quantity".
# these will be used for segmenting the "Sum" column values
posQ = (np.where(df1.Quantity > 0)[0])

# # don't want zero or last index value in posQ for splitting
if posQ[0] == 0:
    posQ = posQ[1:]
if posQ[-1] == len(df)-1:
    posQ = posQ[:-1]

# se (start-end)
# add first and last index to be used for indexing segments of df_sum
se = posQ.copy()
se = np.insert(se, 0, 0)
se = np.append(se, len(df))

starts = se[:-1]
ends = se[1:]

# keep only positive values from the df_sum array.
# this is used with numpy argmin to find first non-positive number
# within segments
only_pos = np.add(np.zeros(len(df)), np.where(df_sum > 0, df_sum, 0))

# split the only_positive array at Quantity locations
segs = np.array(np.split(only_pos, posQ))

# find the indices of the neg numbers within each segment
tgts = [np.argmin(x) for x in segs]

# use the indices to slice each segment and put the result into
# the result array
i = 0
for seg in segs:
    idx = np.arange(starts[i], ends[i])
    np.put(new_sum_array, idx[tgts[i]:], df_sum[idx][tgts[i]:])
    i += 1

# to add a lookback limit for adding consecutive positive df_sums,
# assign an integer value to max_lookback in next line.
# use "None" to ignore any limit
max_lookback = None
if max_lookback is not None:
    tgts = np.clip(tgts, 0, max_lookback)

# add up the values of the positive numbers in the sliced
# df_sum segments
sums = [np.sum(x[:l]) for x, l in zip(segs, tgts)]

# put those totals into the result array at positive "Quality" locations
np.put(new_sum_array, starts, sums)

# add the results to the df as "New Sum"
df1['New Sum'] = new_sum_array

# flip the dataframe back upright
df1 = df1[::-1]
# insert calculated column into original dataframe
df['new sum'] = df1['New Sum']

Результат:

print(df)

    Code  Sum  Quantity  New Sum
0      0  -12         0      -12
1      1   23         0       23
2      2  -10         0      -10
3      3  -12         0      -12
4      4  100         0        0
5      5  102       201      202
6      6   34         0       34
7      7  -34         0      -34
8      8  -23         0      -23
9      9  100         0        0
10    10  100         0        0
11    11  102       300      302

Answer 2

Вы можете сгруппировать кадр данных по кумулятивной сумме количества в обратном порядке, рассчитать все положительные и суммарные значения и присвоить обратно столбцу

df['New sum'] = df[df.Sum.lt(0)]['Sum']
a = df.groupby([df.Quantity.ne(0)[::-1].cumsum()])['Sum'].apply(lambda x:x[x.ge(0)].sum())[::-1]
df['New sum'] = pd.Series(a.values,index=df[df.Quantity.ne(0)].index)

Out:

    Code    Sum Quantity    New sum
0   0   -12 0   NaN
1   1   -23 0   NaN
2   2   -12 0   NaN
3   3   100 0   NaN
4   4   102 201 202.0
5   5   -34 0   NaN
6   6   -23 0   NaN
7   7   100 0   NaN
8   8   100 0   NaN
9   9   102 300 302.0