Put next months start as previous months end pandas - pandas

Question

I have a dataframe in long format (panel data), Each person has a start month along with variables. it looks something like:
Data description
person_id
month_start
Var1
Var2
1
1
0.4
1.4
1
2
0.3
0.131
1
3
0.34
0.434
2
2
0.49
0.949
2
3
0.53
1.53
2
5
0.38
0.738
3
1
1.12
1.34
3
4
1.89
1.02
3
5
0.83
0.27
and I need it to look like:
person_id
month_start
month_end
Var1
Var2
1
1
2
0.4
1.4
1
2
3
0.3
0.131
1
3
4
0.34
0.434
2
2
3
0.49
0.949
2
3
5
0.53
1.53
2
5
6
0.38
0.738
3
1
4
1.12
1.34
3
4
5
1.89
1.02
3
5
6
0.83
0.27
Where month end is the beginning of the next entry for that person.
I was able to make this:
a = pd.DataFrame({'person_id':[1,1,1,2,2,2,3,3,3], 'var1': [0.4, 0.3, 0.34, 0.49, 0.53, 0.38, 1.12, 1.89, 0.83], 'var2': [1.4, 0.131, 0.434, 0.949, 1.53, 0.738, 1.34, 1.02, 0.27], 'month_start': [1,2,3,2,3,5,1,4,5]})
def add_end_date(df_in,object_id, start_col, end_col):
df = df_in.copy()
prev_person_id = -1
prev_index = -1
df[end_col] = [-1]*len(df)
for idx, row in df.iterrows():
p_id = row[object_id]
p_idx = idx
if prev_person_id == p_id:
df.loc[prev_index, end_col] = int(row[start_col])# put in start date as last entries end date
if row[end_col] == -1:
df.loc[idx, end_col] = int(row[start_col]+1)
prev_person_id = p_id
prev_index = p_idx
return df
add_end_date(a, 'person_id', 'month_start', 'month_end')
Is there a better/optimized way to accomplish this?

Answer 1

Try groupby.shift:
df['month_end'] = df.groupby('person_id').month_start.shift(-1)\
.fillna(df.month_start + 1).astype(int)
df
person_id month_start Var1 Var2 month_end
0 1 1 0.40 1.400 2
1 1 2 0.30 0.131 3
2 1 3 0.34 0.434 4
3 2 2 0.49 0.949 3
4 2 3 0.53 1.530 5
5 2 5 0.38 0.738 6
6 3 1 1.12 1.340 4
7 3 4 1.89 1.020 5
8 3 5 0.83 0.270 6