pandas 将连续数字分组求和

看过来

《pandas 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

本例将重点讲解如何对数值按连续段进行分组。pandas 并不提供直接的判断一个序列中有哪些子序列是连续的，需要我们按照一些逻辑技巧来自己实现。本例除了按连续性分组外，还将演示如何完成自定义聚合操作。

数据与需求

我们构造样例数据如下：

import pandas as pd

df = pd.DataFrame({'a': [1, 2, 3, 6, 7, 8, 9, 12],
                   'b': range(8)})
df
'''
    a  b
0   1  0
1   2  1
2   3  2
3   6  3
4   7  4
5   8  5
6   9  6
7  12  7
'''

可以看到，a 列是要判断是否连续的列，b 列是需要聚合计算的列，我们通过观察看到 a 列有以下几个子连续序列组成：

1 到 3
6 到 9
12

共有三段，按以上三段分组后，得到的预期数据结果如下：

# 目标数据
'''
       a   b
0    1-3   3
1    6-9  18
2  12-12   7
'''

思路

此题的重要是对每个值进行标注，将三段数据分别标记为 0、1、2 之类的值，告诉我们哪几个数字在同一个连续序列里。

要对数据进行标记，就需要一个辅助列来完成。我们可以对 a 列求当前值与下一个值的差（位差），如果和上一个连续则位差为 1，不连续位差会大于 1。判断是否大于 1 后得到一个布尔序列，再对这个布尔序列求累积和，由于 False 为 0，只到遇到 True （为 1）时标记值才能增加一。

这样我们就巧妙标记出不同的连续分组，后边的问题就简单了。

最后我们还会尝试一个非常巧妙的利用 rank() 排序序号来分组的办法。

代码

对 a 列求位差：

df.a.diff(1)
'''
0    NaN
1    1.0
2    1.0
3    3.0
4    1.0
5    1.0
6    1.0
7    3.0
Name: a, dtype: float64
'''

df.a.diff(1).fillna(1)>1
'''
0    False
1    False
2    False
3     True
4    False
5    False
6    False
7     True
Name: a, dtype: bool
'''

(df.a.diff(1).fillna(1)>1).cumsum()
'''
0    0
1    0
2    0
3    1
4    1
5    1
6    1
7    2
Name: a, dtype: int64
'''

将以上结果增加到虚拟辅助列：

df.assign(a_1=(df.a.diff(1).fillna(1)>1).cumsum())
'''
    a  b  a_1
0   1  0    0
1   2  1    0
2   3  2    0
3   6  3    1
4   7  4    1
5   8  5    1
6   9  6    1
7  12  7    2
'''

接着，按这个辅助列进行分组聚合，完成后再删除这个辅助列：

(
    df.assign(a_1=(df.a.diff(1).fillna(1)>1).cumsum())
    .groupby('a_1', as_index=False)
    .agg({'a': lambda x: f'{x.min()}-{x.max()}',
          'b': sum})
    .drop('a_1', axis=1)
)
'''
       a   b
0    1-3   3
1    6-9  18
2  12-12   7
'''

rank 分组

还有一种比较巧妙的标记连续的方法是用此列减去此列的 rank() 排序值：

df.assign(rank=df.a.rank(),
          mark=df.a - df.a.rank())
'''
    a  b  rank  mark
0   1  0   1.0   0.0
1   2  1   2.0   0.0
2   3  2   3.0   0.0
3   6  3   4.0   2.0
4   7  4   5.0   2.0
5   8  5   6.0   2.0
6   9  6   7.0   2.0
7  12  7   8.0   4.0
'''

我们可以看到，如果是连续的值，它们的 rank() 序号也是连续的，它减去序号的值是相同的，这样就将它们标记在同一个组。

我们可以直接利用这个标记进行分组：

(
    df.groupby(df.a - df.a.rank(), as_index=False)
    .agg({'a': lambda x: f'{x.min()}-{x.max()}',
          'b': sum})
)
'''
       a   b
0    1-3   3
1    6-9  18
2  12-12   7
'''

如果将一组中的最大和最小值单独列为一列，代码则修改为：

(
    df.groupby(df.a - df.a.rank(), as_index=False)
    .agg(a_min=('a', min),
        b_max=('a', max),
        b=('b', sum))
)
'''
   a_min  b_max   b
0      1      3   3
1      6      9  18
2     12     12   7
'''

这样就得到了我们想要的结果。

（完）

pandas 将连续数字分组求和

数据与需求

思路

代码

rank 分组

相关内容