Как создать инкремент var из первого значения группы данных?

Question

У меня есть дата:

data=[[0,1,5],
      [0,1,6],
      [0,0,8],
      [0,0,10],
      [0,1,12],
      [0,0,14],
      [0,1,16],
      [0,1,18],    
      [1,0,2],
      [1,1,0],
      [1,0,1],
      [1,0,2]]

df = pd.DataFrame(data,columns=['KEY','COND','VAL'])

Для RES1 я хочу создать переменную счетчика RES, где COND == 1. Значение RES для первого KEY группы остается таким же, как VAL (Могу ли я использовать cumcount () каким-то образом).

Для RES2, тогда я просто хочу заполнить пропущенные значения в качестве предыдущего значения. (df.fillna(method='ffill')), я думаю ...

    KEY  COND    VAL RES1  RES2
0     0     1    5      5     5
1     0     1    6      6     6
2     0     0    8            6
3     0     0   10            6
4     0     1   12      7     7
5     0     0   14            7
6     0     1   16      8     8
7     0     1   18      9     9
8     1     0    2      2     2
9     1     1    0      3     3
10    1     0    1            3
11    1     0    2            3

Цель - найти векторизованное решение, наиболее оптимальное для миллионов строк.

Answer 1

IIUC

con=(df.COND==1)|(df.index.isin(df.drop_duplicates('KEY').index))
df['res1']=df.groupby('KEY').VAL.transform('first')+
           df.groupby('KEY').COND.cumsum()[con]-
           df.groupby('KEY').COND.transform('first')
df['res2']=df.res1.ffill()
df
   Out[148]: 
    KEY  COND  VAL  res1  res2
0     0     1    5   5.0   5.0
1     0     1    6   6.0   6.0
2     0     0    8   NaN   6.0
3     0     0   10   NaN   6.0
4     0     1   12   7.0   7.0
5     0     0   14   NaN   7.0
6     0     1   16   8.0   8.0
7     0     1   18   9.0   9.0
8     1     0    2   2.0   2.0
9     1     1    0   3.0   3.0
10    1     0    1   NaN   3.0
11    1     0    2   NaN   3.0

Answer 2

Вы хотите:

s = (df[df.KEY.duplicated()]          # Ignore first row in each KEY group
        .groupby('KEY').COND.cumsum() # Counter within KEY
        .add(df.groupby('KEY').VAL.transform('first')) # Add first value per KEY
        .where(df.COND.eq(1))                          # Set only where COND == 1
        .add(df.loc[~df.KEY.duplicated(), 'VAL'], fill_value=0) # Set 1st row by KEY
     )
df['RES1'] = s
df['RES2'] = df['RES1'].ffill()

---

    KEY  COND  VAL  RES1  RES2
0     0     1    5   5.0   5.0
1     0     1    6   6.0   6.0
2     0     0    8   NaN   6.0
3     0     0   10   NaN   6.0
4     0     1   12   7.0   7.0
5     0     0   14   NaN   7.0
6     0     1   16   8.0   8.0
7     0     1   18   9.0   9.0
8     1     0    2   2.0   2.0
9     1     1    0   3.0   3.0
10    1     0    1   NaN   3.0
11    1     0    2   NaN   3.0