Group rows partially [Python] [Pandas]

Question

Good morning everyone.

I have the following data:

import pandas as pd

info = {
'states': [-1, -1, -1, 1, 1, -1, 0, 1, 1, 1],
'values': [34, 29, 28, 30, 35, 33, 33, 36, 40, 41] }

df = pd.DataFrame(data=info)

print(df)

>>> 
    states   values
0       -1       34
1       -1       29
2       -1       28
3        1       30
4        1       35
5       -1       33
6        0       33
7        1       36
8        1       40
9        1       41

I need to group the data using PANDAS (and/or higher order functions) (already did the exercise using for loops), I need to group the data having the "states" column as a guide. But the grouping should not be of all the data, I only need to group the data that is neighboring... as follows:

Initial DataFrame:

    states   values
0       -1       34 ┐
1       -1       29 │    Group this part (states = -1)
2       -1       28 ┘
3        1       30 ┐    Group this part (states =  1)
4        1       35 ┘
5       -1       33     'Group' this part (states = -1)
6        0       33     'Group' this part (states =  0)
7        1       36 ┐
8        1       40 │    Group this part (states =  1)
9        1       41 ┘

It would result in a DataFrame, with a grouping by segments (from the "states" column) and in another column the sum of the data (from the "values" column).

Expected DataFrame:

    states   values
0       -1       91     (values=34+29+28)
1        1       65     (values=30+35)
2       -1       33
3        0       33
4        1       117    (values=36+40+41)

You who are more versed in these issues, perhaps you can help me perform this operation.

Thank you so much!

Answer 1

You can do this by making use of shift to create different groups based on consecutive values of the states column, after which you can use groupby to add the values in the values column:

(
    df
    .assign(
        # create groups by checking if value is different from previous value
        group = lambda x: (x["states"] != x["states"].shift()).cumsum()
    )
    # group by the states and group columns
    .groupby(["states", "group"])
    # sum the values in the values column
    .agg({"values": "sum"})
    .reset_index()
    # select the columns needed in the final output
    .loc[:, ["group", "values"]]
)

Answer 2

It's not the problem about grouping the elements, it is the problem related to consecutive elements. The approach should be of using consecutive iterations.

Here I've saved all the results in a list. Don't use any dictionary to store results as its key values should be unique.

Your desired output ☝️