Efficiently finding consecutive streaks in a pandas DataFrame column?

Viewed 610

I have a DataFrame similar to the below:, and I want to add a Streak column to it (see example below):

Date         Home_Team    Away_Team    Winner      Streak

2005-08-06       A            G           A           0
2005-08-06       B            H           H           0
2005-08-06       C            I           C           0
2005-08-06       D            J           J           0
2005-08-06       E            K           K           0
2005-08-06       F            L           F           0
2005-08-13       A            B           A           1           
2005-08-13       C            D           D           1           
2005-08-13       E            F           F           0        
2005-08-13       G            H           H           0
2005-08-13       I            J           J           0
2005-08-13       K            L           K           1
2005-08-20       B            C           B           0
2005-08-20       A            D           A           2
2005-08-20       G            K           K           0
2005-08-20       I            E           E           0
2005-08-20       F            H           F           2
2005-08-20       J            L           J           2
2005-08-27       A            H           A           3
2005-08-27       B            F           B           1
2005-08-27       J            C           C           3           
2005-08-27       D            E           D           0
2005-08-27       I            K           K           0
2005-08-27       L            G           G           0
2005-09-05       B            A           A           2
2005-09-05       D            C           D           1
2005-09-05       F            E           F           0
2005-09-05       H            G           H           0
2005-09-05       J            I           I           0
2005-09-05       K            L           K           4

The DataFrame is approximately 200k rows going from 2005 to 2020.

Now, what I am trying to do is find the number of consecutive games the Home Team has won PRIOR to the date in in the Date column in the DataFrame. I have a solution, but it is too slow, see below:

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)

How can I speed this up?

4 Answers

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)

Can't think of a pandas solution, but you can assign a group number using ngroup and then use defaultdict to create groups so you can lookup the accumulate results:

from collections import defaultdict

d = defaultdict(lambda: defaultdict(int))

df["group"] = df.groupby("Date").ngroup()

for a, b in zip(df["Winner"], df["group"]):
    d[b][a] = 1+d.get(b-1,{}).get(a, 0)

df["Streak"] = [d.get(y-1, {}).get(x, 0) for x, y in zip(df["Home_Team"], df["group"])]

print (df.drop("group", 1))

          Date Home_Team Away_Team Winner  Streak
0   2005-08-06         A         G      A       0
1   2005-08-06         B         H      H       0
2   2005-08-06         C         I      C       0
3   2005-08-06         D         J      J       0
4   2005-08-06         E         K      K       0
5   2005-08-06         F         L      F       0
6   2005-08-13         A         B      A       1
7   2005-08-13         C         D      D       1
8   2005-08-13         E         F      F       0
9   2005-08-13         G         H      H       0
10  2005-08-13         I         J      J       0
11  2005-08-13         K         L      K       1
12  2005-08-20         B         C      B       0
13  2005-08-20         A         D      A       2
14  2005-08-20         G         K      K       0
15  2005-08-20         I         E      E       0
16  2005-08-20         F         H      F       2
17  2005-08-20         J         L      J       2
18  2005-08-27         A         H      A       3
19  2005-08-27         B         F      B       1
20  2005-08-27         J         C      C       3
21  2005-08-27         D         E      D       0
22  2005-08-27         I         K      K       0
23  2005-08-27         L         G      G       0
24  2005-09-05         B         A      A       2
25  2005-09-05         D         C      D       1
26  2005-09-05         F         E      F       0
27  2005-09-05         H         G      H       0
28  2005-09-05         J         I      I       0
29  2005-09-05         K         L      K       4

FIXES IN PROGRESS!

Here is probably the simplest approach -

def get_streak(l,m,n):
    wins = np.roll(np.cumsum([1 if i==n else 0 for i in l]),1)
    wins[0]=0
    filts = np.array([1 if i==n else 0 for i in m])
    mul = np.multiply(wins, filts)
    return mul


streaks = np.zeros((30,)).astype(int)
l = list(df['Winner'])
m = list(df['Home_Team'])

for i in df['Winner'].unique():
    streaks += get_streak(l,m,i)
    
df['streaks'] = streaks
          Date Home_Team Away_Team Winner  streaks
0   2005-08-06         A         G      A        0
1   2005-08-06         B         H      H        0
2   2005-08-06         C         I      C        0
3   2005-08-06         D         J      J        0
4   2005-08-06         E         K      K        0
5   2005-08-06         F         L      F        0
6   2005-08-13         A         B      A        1
7   2005-08-13         C         D      D        1
8   2005-08-13         E         F      F        0
9   2005-08-13         G         H      H        0
10  2005-08-13         I         J      J        0
11  2005-08-13         K         L      K        1
12  2005-08-20         B         C      B        0
13  2005-08-20         A         D      A        2
14  2005-08-20         G         K      K        0
15  2005-08-20         I         E      E        0
16  2005-08-20         F         H      F        2
17  2005-08-20         J         L      J        2
18  2005-08-27         A         H      A        3
19  2005-08-27         B         F      B        1
20  2005-08-27         J         C      C        3
21  2005-08-27         D         E      D        1
22  2005-08-27         I         K      K        0
23  2005-08-27         L         G      G        0
24  2005-09-05         B         A      A        2
25  2005-09-05         D         C      D        2
26  2005-09-05         F         E      F        3
27  2005-09-05         H         G      H        2
28  2005-09-05         J         I      I        3
29  2005-09-05         K         L      K        4

It's pretty straightforward -

  1. You take a cumulative sum of the wins of a given team and shift them by 1.
  2. Then you take the dot product of those with the instances where they were the home team. Save that to a vector called streak
  3. You loop through all the unique home teams and calculate the sum of their streaks.
  4. Done!

A more intuitive sense of how the function works can be seen with some print statements -

def get_streak(l,m,n):
    wins = np.roll(np.cumsum([1 if i==n else 0 for i in l]),1)
    wins[0]=0
    print('wins:',wins)
    filts = np.array([1 if i==n else 0 for i in m])
    print('home:',filts)
    mul = np.multiply(wins, filts)
    print('strk:', mul)
    return mul

streak_A = get_streak(l,m,'A')
wins: [0 1 1 1 1 1 1 2 2 2 2 2 2 2 3 3 3 3 3 4 4 4 4 4 4 5 5 5 5 5]
home: [1 0 0 0 0 0 1 0 0 0 0 0 0 1 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0]
strk: [0 0 0 0 0 0 1 0 0 0 0 0 0 2 0 0 0 0 3 0 0 0 0 0 0 0 0 0 0 0]

The element-wise sum of all the streaks is what you are looking for.


Benchmark (seems to be the fastest among all other answers) -

529 µs ± 20.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Related