有效地在 pandas DataFrame 列中找到连续条纹？答案

【问题标题】：Efficiently finding consecutive streaks in a pandas DataFrame column?有效地在 pandas DataFrame 列中找到连续条纹？
【发布时间】：2020-12-19 15:45:20
【问题描述】：

我有一个类似于下面的 DataFrame：，我想向它添加一个 Streak 列（参见下面的示例）：

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

从 2005 年到 2020 年，DataFrame 大约有 20 万行。

现在，我要做的是在 DataFrame 的 Date 列中找到主队在该日期之前赢得的连续比赛数。我有一个解决办法，但是太慢了，见下文：

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)

如何加快速度？

【问题讨论】：

所以你想按主队分组并计算连续获胜？
如果A 作为客队获胜会怎样？如果输了怎么办？这会继续/结束连胜吗？还是信息丢失了？
它继续/结束连胜。主场/客场与连续上垒无关，我只想要主队的连续上垒。
我明白了。所以你可以看到跳跃吗？您如何知道以客队结束的球队的连胜纪录？您可以随时为客场连胜添加一列...
有道理。我只是想在回答之前确保我完全理解了这个问题。最后一个问题：您对 numpy 而不是 pandas 解决方案还满意吗？

标签： python pandas numpy dataframe

【解决方案1】：

我将在这里介绍一个基于 numpy 的解决方案。首先是因为我对 pandas 不是很熟悉，也不想进行研究，其次是因为 numpy 解决方案无论如何都应该可以正常工作。

让我们先来看看特定团队会发生什么。您的目标是根据参加的比赛顺序找出球队连续获胜的次数。我将删除日期列并将您的数据转换为初学者的 numpy 数组：

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']])

您不需要日期，因为您只关心谁玩过，即使他们在一天内玩了多次。那么让我们来看看刚刚的团队A：

A_played = np.flatnonzero((x[:, :2] == 'A').any(axis=1))
A_won = x[A_played, -1] == 'A'

A_played 是一个索引数组，其元素数与x 中的行数相同。 A_won 是一个具有与np.count_nonzero(A_played) 一样多的元素的掩码；即A参与的游戏数。

找出条纹的大小是一个经过充分讨论的问题：

streaks = np.diff(np.flatnonzero(np.diff(np.r_[False, A_won, False])))[::2]

您计算掩码值切换的每对索引之间的差异。 False 的额外填充确保您知道掩码切换的方式。您正在寻找的是基于此计算但需要更多细节，因为您想要累积和，但在每次运行后重置。您可以通过在运行后立即将数据的值设置为否定的运行长度来做到这一点：

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

现在你有一个可修剪的数组，它的累积和可以直接分配给输出列：

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]

现在您可以遍历所有团队（与比赛总数相比，这应该是一个相当小的数字）：

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)

【讨论】：

这个解决方案似乎几乎是正确的，但我认为它错过了一个事实，即我想查看在玩游戏之前赢得的连胜，这个解决方案似乎将1 添加到与比赛在同一行的连胜（如果他们赢了），什么时候应该将其添加到下一场比赛的行中的连胜中
所以第一次球队比赛时，连胜应该被初始化为0，从那里开始，连胜取决于之前的比赛（我们正在考虑的当前比赛之前的所有比赛） .
@theman 谢谢你的收获。我已经适当地调整了合同。我知道您的数据不存在这种情况，但由于我们删除了时间索引，因此无论团队每天玩多少场比赛，此解决方案都有效。
这也不太行。似乎一支球队必须至少连续两次获胜才能出现在连胜纪录中。此外，当一支球队输球时，该行的连胜纪录似乎被设置为零，而不是下一场比赛
@theman。这就是我在移动设备上发生的情况 :) 当我可以在桌面上测试时我会修复

【解决方案2】：

优雅的方式：

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

据了解，一支球队每天比赛不超过一次

时代

%%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)

【讨论】：

这似乎对我也不起作用；我对条纹 (100+) 的看法很疯狂，我知道这不是真的。
一支球队一天可以打超过一次吗？
不，每天最多一次

【解决方案3】：

想不出pandas 解决方案，但您可以使用ngroup 分配组号，然后使用defaultdict 创建组，以便查找累积结果：

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

【讨论】：

这似乎对我不起作用。我在 Streak 列中得到的唯一值是 0、1 和 2，在我较大的 DataFrame 中，我知道有比 2 长得多的条纹
我现在添加了更多数据
它工作得很好。您的问题出在其他地方，可能是数据不一致。
嗯，有时存在关联的事实是否会影响您的解决方案？如在Winner 列中，偶尔会有："tie"。我现在才注意到这一点，抱歉。
它不应该改变任何东西。将有一个名为“tie”的新组，但不会影响其他组。您确定团队的名称在您的数据中保持不变，没有错别字吗？

【解决方案4】：

修复中！

这可能是最简单的方法 -

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

这很简单-

您将给定团队的胜利累积总和并将它们移动 1。
然后，您将那些与他们是主队的情况相乘的点积。将其保存到名为 streak 的向量中
您遍历所有独特的主队并计算他们的连胜纪录的总和。
完成！

通过一些打印语句可以更直观地了解函数的工作原理 -

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]

所有条纹的元素总和就是您要寻找的。p>

基准测试（似乎是所有其他答案中最快的）-

每个循环 529 µs ± 20.6 µs（7 次运行的平均值 ± 标准偏差，每次 1000 个循环）

【讨论】：