CODEWITHSUNDEEP
pythonpandas
exp1orer
exp1orer

Reputation: 12019

How to add "order within group" column in pandas?

Take the following dataframe:

import pandas as pd
df = pd.DataFrame({'group_name': ['A','A','A','B','B','B'],
                   'timestamp': [4,6,1000,5,8,100],
                   'condition': [True,True,False,True,False,True]})

I want to add two columns:

  1. The row's order within its group
  2. rolling sum of the condition column within each group

I know I can do it with a custom apply, but I'm wondering if anyone has any fun ideas? (Also this is slow when there are many groups.) Here's one solution: 

def range_within_group(input_df):
    df_to_return = input_df.copy()
    df_to_return = df_to_return.sort('timestamp')
    df_to_return['order_within_group'] = range(len(df_to_return))
    df_to_return['rolling_sum_of_condition'] = df_to_return.condition.cumsum()
    return df_to_return

df.groupby('group_name').apply(range_within_group).reset_index(drop=True)

Upvotes: 5

Views: 5312

Answers (1)

behzad.nouri
behzad.nouri

Reputation: 77951

GroupBy.cumcount does:

Number each item in each group from 0 to the length of that group - 1.

so simply:

>>> gr = df.sort('timestamp').groupby('group_name')
>>> df['order_within_group'] = gr.cumcount()
>>> df['rolling_sum_of_condition'] = gr['condition'].cumsum()

On pandas >= 0.2
df.sort() is not valid anymore, you have to use df.sort_values()

>>> gr = df.sort_values('timestamp').groupby('group_name')
>>> df['order_within_group'] = gr.cumcount()
>>> df['rolling_sum_of_condition'] = gr['condition'].cumsum()

Upvotes: 7

Related Questions