Iterating over a dataframe twice: which is the ideal way? - pandas

Question

I am trying to create a dataframe for Sankey chart in Power BI which needs source and destination like this.
id
Source
Destination
1
Starting a
next point b
1
next point b
final point c
1
final point c
end
2
Starting a
next point b
2
next point b
3
Starting a
next point b
3
next point b
final point c
3
final point c
end
I have a dataframe like this:
ID
flow
1
Starting a
1
next point b
1
final point c
2
Starting a
2
next point b
3
Starting a
3
next point b
3
final point c
I tried doing by iterating over the dataframe twice like below:
for index, row in df.iterrows():
for j, r in df.iterrows():
if row['ID'] == r['ID']:
if (index + 1 == j) & ("final point c" not in row['flow']):
df['Destination'][index] = df['flow'][j]
elif "final point c" in row['flow']:
df['Destination'][index] = 'End of flow'
Since it is iterating over the same dataframe twice, when the records are huge, it is taking a lot of time to process.
Is there any better way to do this? I tried looking at the all similar questions, but couldn't find anything that relates to my question.

Answer 1

You could use groupby+shift and a bit of masking:
end = df['flow'].str.startswith('final point')
df2 = (df.assign(destination=df.groupby('ID')['flow'].shift(-1)
.mask(end, end.map({True: 'end'}))
)
.rename(columns={'flow': 'source'})
)
output:
ID source destination
0 1 Starting a next point b
1 1 next point b final point c
2 1 final point c end
3 2 Starting a next point b
4 2 next point b NaN
5 3 Starting a next point b
6 3 next point b final point c
7 3 final point c end
Alternative with combine_first to fill the NaNs:
end = df['flow'].str.startswith('final point').map({True: 'end', False: ''})
df2 = (df.assign(destination=df.groupby('ID')['flow'].shift(-1).combine_first(end))
.rename(columns={'flow': 'source'})
)
output:
ID source destination
0 1 Starting a next point b
1 1 next point b final point c
2 1 final point c end
3 2 Starting a next point b
4 2 next point b
5 3 Starting a next point b
6 3 next point b final point c
7 3 final point c end