Add values to a group and recalculate in Python-CodePudding

I have a dataset where I would like to add certain values to a certain row and recalculate the values of other specific rows.

For example. For 'aa' q222, I would like to add the value 500 to the 'pwr' column. This will then change the 'avail' values for q222 onward.

Data

id  date    pwr     used    avail
aa  q122    1000    727.2   272.8
aa  q222    1000    975.7   24.3
aa  q322    1000    1092.7  -92.7
aa  q422    1000    1196.7  -196.7
bb  q122    1000    0       1000
bb  q222    1000    117.3   882.7

Desired

id  date    pwr     used    avail
aa  q122    1000    727.2   272.8
aa  q222    1500    975.7   524.3
aa  q322    1500    1092.7  407.3
aa  q422    1500    1196.7  303.3
bb  q122    1000    0       1000
bb  q222    1100    117.3   982.7

Logic

added 500 to the 'pwr' column for aa q222 onward
added 100 to the 'pwr' column for bb q222 onward

Doing

I am thinking I can index the row, but not sure exactly how to do this. I am still researching. Any suggestion is appreciated.

df.loc[df.index['q222], 'aa'] = 500

CodePudding user response：

Here we go:

# row q222 onward, within `id` group
mask = df['date'].eq('q222').groupby(df['id']).cumsum() >= 1

df.loc[mask, 'avail']  = df['id'].map({'aa':500, 'bb':100})

Output:

   id  date   pwr    used   avail
0  aa  q122  1000   727.2   272.8
1  aa  q222  1000   975.7   524.3
2  aa  q322  1000  1092.7   407.3
3  aa  q422  1000  1196.7   303.3
4  bb  q122  1000     0.0  1000.0
5  bb  q222  1000   117.3   982.7

Note if your date can be ordered as the name suggests, and the onward reflects the ordering, then mask can simply be:

mask = df['date'] >= 'q222'