Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import pandas as pd
df = pd.DataFrame({'agent_id': [1, 1, 1, 1, 2, 2, 2, 2],
'date': ["2021-01-01", "2021-04-01", "2021-05-01", "2021-06-01", "2021-01-01", "2021-02-01", "2021-03-01", "2021-06-01"],
'txn_amount': [100, 200, 100, 200, 100, 200, 100, 200],
'txn_status': ["Failure", "Success", "Failure", "Success", "Failure", "Success", "Failure", "Success"]})
df['date'] = pd.to_datetime(df['date'], format='%Y-%m-%d')
df1 = df.set_index('date')
df1['month'] = df1.index.month
df1['year'] = df1.index.year
df1.groupby([df1['agent_id'] , df1.year]).apply(is_at_least_three_consec)
def is_at_least_three_consec(grouped_df):
"""
Function to check for 3 consecutive months
Input: Grouped df
Output: Subset of the df which has continuous month values
"""
month_diff = grouped_df['month'].diff().values.tolist()
consec_count = 0
for index , val in enumerate(month_diff):
if index != 0 and val == 1:
consec_count += 1
if consec_count == 2:
return grouped_df[index-2: index+1]
else:
consec_count = 0
return None
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment