Efficiently finding consecutive streaks in a pandas DataFrame column?
I will present a numpy-based solution here. Firstly because I am not very familiar with pandas and don't feel like doing the research, and secondly because a numpy solution should work just fine regardless.
Let's take a look at what happens to one given team first. Your goal is to find the number of consecutive wins for a team based on the sequence of games it participated in. I will drop the date column and turn your data into a numpy array for starters:
x = np.array([
['A', 'G', 'A'],
['B', 'H', 'H'],
['C', 'I', 'C'],
['D', 'J', 'J'],
['E', 'K', 'K'],
['F', 'L', 'F'],
['A', 'B', 'A'],
['C', 'D', 'D'],
['E', 'F', 'F'],
['G', 'H', 'H'],
['I', 'J', 'J'],
['K', 'L', 'K'],
['B', 'C', 'B'],
['A', 'D', 'A'],
['G', 'K', 'K'],
['I', 'E', 'E'],
['F', 'H', 'F'],
['J', 'L', 'J']])
You don't need the date because all you care about is who played, even if they did it multiple times in one day. So let's take a look at just team A:
A_played = np.flatnonzero((x[:, :2] == 'A').any(axis=1))
A_won = x[A_played, -1] == 'A'
A_played is an index array with the same number of elements as there are rows in x. A_won is a mask that has as many elements as np.count_nonzero(A_played); i.e., the number of games A participated in.
Finding the sizes of the streaks is a fairly well hashed out problem:
streaks = np.diff(np.flatnonzero(np.diff(np.r_[False, A_won, False])))[::2]
You compute the differences between each pair of indices where the value of the mask switches. The extra padding with False ensures that you know which way the mask is switching. What you are looking for is based on this computation but requires a bit more detail, since you want the cumulative sum, but reset after each run. You can do that by setting the value of the data to the negated run length immediately after the run:
wins = np.r_[0, A_won, 0] # Notice the int dtype here
switch_indices = np.flatnonzero(np.diff(wins)) + 1
streaks = np.diff(switch_indices)[::2]
wins[switch_indices[1::2]] = -streaks
Now you have a trimmable array whose cumulative sum can be assigned directly to the output columns:
streak_counts = np.cumsum(wins[:-2])
output = np.zeros((x.shape[0], 2), dtype=int)
# Home streak
home_mask = x[A_played, 0] == 'A'
output[A_played[home_mask], 0] = streak_counts[home_mask]
# Away streak
away_mask = ~home_mask
output[A_played[away_mask], 1] = streak_counts[away_mask]
Now you can loop over all teams (which should be a fairly small number compared to the total number of games):
def process_team(data, team, output):
played = np.flatnonzero((data[:, :2] == team).any(axis=1))
won = data[played, -1] == team
wins = np.r_[0, won, 0]
switch_indices = np.flatnonzero(np.diff(wins)) + 1
streaks = np.diff(switch_indices)[::2]
wins[switch_indices[1::2]] = -streaks
streak_counts = np.cumsum(wins[:-2])
home_mask = data[played, 0] == team
away_mask = ~home_mask
output[played[home_mask], 0] = streak_counts[home_mask]
output[played[away_mask], 1] = streak_counts[away_mask]
output = np.empty((x.shape[0], 2), dtype=int)
# Assume every team has been home team at least once.
# If not, x[:, :2].ravel() copies the data and np.unique(x[:, :2]) does too
for team in set(x[:, 0]):
process_team(x, team, output)
Elegant way:
new_df = (df.reset_index()
.melt(['index', 'Date', 'Winner'])
.assign(win=lambda x: x['value'].eq(x.Winner))
.sort_values('Date')
.assign(cum_wins=lambda x: x.groupby('value')['win'].cumsum())
.assign(cum_wins_prev=lambda x: x.groupby('value')['cum_wins'].shift(fill_value=0))
.pivot_table(index='index', values='cum_wins_prev', columns='variable')
.add_prefix('Streak_')
)
print(new_df)
variable Streak_Away_Team Streak_Home_Team
index
0 0.0 0.0
1 0.0 0.0
2 0.0 0.0
3 0.0 0.0
4 0.0 0.0
5 0.0 0.0
6 0.0 1.0
7 0.0 1.0
8 1.0 0.0
9 1.0 0.0
10 1.0 0.0
11 0.0 1.0
12 1.0 0.0
13 1.0 2.0
14 2.0 0.0
15 0.0 0.0
16 2.0 2.0
17 0.0 2.0
#new_df = df.assign(**new_df) #you could use join or assign
new_df = df.join(new_df)
print(new_df)
Date Home_Team Away_Team Winner Streak_Away_Team Streak_Home_Team
0 2005-08-06 A G A 0.0 0.0
1 2005-08-06 B H H 0.0 0.0
2 2005-08-06 C I C 0.0 0.0
3 2005-08-06 D J J 0.0 0.0
4 2005-08-06 E K K 0.0 0.0
5 2005-08-06 F L F 0.0 0.0
6 2005-08-13 A B A 0.0 1.0
7 2005-08-13 C D D 0.0 1.0
8 2005-08-13 E F F 1.0 0.0
9 2005-08-13 G H H 1.0 0.0
10 2005-08-13 I J J 1.0 0.0
11 2005-08-13 K L K 0.0 1.0
12 2005-08-20 B C B 1.0 0.0
13 2005-08-20 A D A 1.0 2.0
14 2005-08-20 G K K 2.0 0.0
15 2005-08-20 I E E 0.0 0.0
16 2005-08-20 F H F 2.0 2.0
17 2005-08-20 J L J 0.0 2.0
it is understood that a team does not play more than once a day
Times
%%timeit
df["Streak"] = 0
def home_streak(x): # x is a row of the DataFrame
"""Keep track of a team's winstreak"""
home_team = x["Home_Team"]
date = x["Date"]
# all previous matches for the home team
home_df = df[(df["Home_Team"] == home_team) | (df["Away_Team"] == home_team)]
home_df = home_df[home_df["Date"] < date].sort_values(by="Date", ascending=False).reset_index()
if len(home_df.index) == 0: # no previous matches for that team, so start streak at 0
return 0
elif home_df.iloc[0]["Winner"] != home_team: # lost the last match
return 0
else: # they won the last game
winners = home_df["Winner"]
streak = 0
for i in winners.index:
if home_df.iloc[i]["Winner"] == home_team:
streak += 1
else: # they lost, return the streak
return streak
df["Streak"] = df.apply(lambda x: home_streak(x), axis = 1)
66.2 ms ± 9.54 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%%timeit
new_df = (df.reset_index()
.melt(['index', 'Date', 'Winner'])
.assign(win=lambda x: x['value'].eq(x.Winner))
.sort_values('Date')
.assign(cum_wins=lambda x: x.groupby('value')['win'].cumsum())
.assign(cum_wins_prev=lambda x: x.groupby('value')['cum_wins'].shift(fill_value=0))
.pivot_table(index='index', values='cum_wins_prev', columns='variable')
.add_prefix('Streak_')
)
new_df=df.assign(**new_df)
29.5 ms ± 2.97 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)