Python pandas dataframe sum numbers until they change

Viewed 106

I have a simple question and tried everything I know about dataframes but cannot get it right.

This is the code.

import pandas as pd

df = pd.DataFrame([1,1,1,-1,1,-1,-1,-1,0,1,1,1,-1], columns=['A'])
print(df)

Output.

    A
0   1
1   1
2   1
3  -1
4   1
5  -1
6  -1
7  -1
8   0
9   1
10  1
11  1
12 -1

I want a new column which sums the same numbers until they change. So the output should look like this.

    A  B
0   1  1
1   1  2
2   1  3
3  -1 -1
4   1  1
5  -1 -1
6  -1 -2
7  -1 -3
8   0  0
9   1  1
10  1  2
11  1  3
12 -1 -1

The last thing I came up with is this.

df['B'] = df.index
df.loc[df['A'].ne(df['A'].shift(1)), 'D'] = df['B']
df['D'] = df['D'].ffill().astype(int)
df['E'] = df['A'].iloc[df['D']:df['B']].sum()

But it gives an error.

TypeError: cannot do positional indexing on RangeIndex with these indexers [0 0

There must be an easy way to achieve this?

1 Answers

Try:

df["B"] = df.A.groupby((df.A != df.A.shift()).cumsum()).cumsum()
print(df)

Prints:

    A  B
0   1  1
1   1  2
2   1  3
3  -1 -1
4   1  1
5  -1 -1
6  -1 -2
7  -1 -3
8   0  0
9   1  1
10  1  2
11  1  3
12 -1 -1
Related