aashish-chaubey · June 30, 2021 22:13 · Jun 30, 2021 · Jun 30, 2021
diff --git a/consective_months.py b/consective_months.py
@@ -28,4 +28,4 @@ def is_at_least_three_consec(grouped_df):
                     return grouped_df[index-2: index+1]
         else:
             consec_count = 0
-    return -1
+    return None
diff --git a/consective_months.py b/consective_months.py
@@ -0,0 +1,31 @@
+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 -1