Python: adding observations based on an index restart in pandas dataframe

April 5, 2023

I have a dataframe that looks like this:

df = pd.DataFrame({'index': [ 1, 2, 3, 4, 5, 6, 1, 2, 3, 4,5,6 ],
         'data': [  1.5, 0.22, 0.323, 4.4, 5.62, 0.56, 1.32, 2.1, 3.09, 4,5.3,0.6]})

I would like to have a new month when the index count starts again at 1. The month allocation does not follow a pattern, but I created a list that stores the months I want to have attributed to each restart of the count.

What I want is something like this dataframe:

months =["April", "June"]

df = pd.DataFrame({'index': [ 1, 2, 3, 4, 5, 6, 1, 2, 3, 4,5,6 ], 
          'data': [  1.5, 0.22, 0.323, 4.4, 5.62, 0.56, 1.32, 2.1, 3.09, 4,5.3,0.6],
         'Month': ['April', 'April', 'April', 'April', 'April', 'April', 'June', 'June', 'June', 'June', 'June', 'June'] })

>Solution :

you can iterate over the index column:

import pandas as pd

df = pd.DataFrame({'index': [ 1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6 ],
                   'data': [ 1.5, 0.22, 0.323, 4.4, 5.62, 0.56, 1.32, 2.1, 3.09, 4, 5.3, 0.6]})

months = ["April", "June"]
month_index = 0
month_dict = {}

for i in range(len(df)):
    if df.loc[i, 'index'] == 1:
        month_index = (month_index + 1) % len(months)
        month_dict[1] = months[month_index]
    month_dict[df.loc[i, 'index']] = months[month_index]

df['Month'] = df['index'].map(month_dict)

print(df)