pandas 实现类似 SQL row_number() 分组行号

看过来

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

在一些数据处理场景下，很多分析师都喜欢用 pandas 替代 SQL 的相关操作。今天我们就来用 pandas 实现 SQL 中最常见的窗口函数——row_number()，为各组增加一个行号。

数据与需求

我们先构造如下数据：

import pandas as pd
import random

random.seed(666)

df = pd.DataFrame({'a': [1,2,2,2,1,2,1,1],
                   'b': random.choices(range(10), k=8)
                  })

df
'''
   a  b
0  1  4
1  2  9
2  2  4
3  2  5
4  1  8
5  2  5
6  1  7
7  1  1
'''

a 列为分组，b 列为值（本例中为非必须列），需求为增加一个 row_number 列，显示按 a 分组的行号，即所在行的顺序。

思路

这个需求首先要分组，然后我们有两个思路，一个是直接使用分组对象的 cumcount() 方法，它返回每组中的每个项目的编号，从0到该组的长度减1。相当于对每组做如下操作：

self.apply(lambda x: pd.Series(np.arange(len(x)), x.index))

还有一个方法是使用分组对象的 transform() 方法调用 Series 的 rank() 方法。

接下来，我们都来实现一下。

代码

首先第一个思路，分组并使用 cumcount()：

df.groupby('a').cumcount()
'''
0    0
1    0
2    1
3    2
4    1
5    3
6    2
7    3
dtype: int64
'''

由于行号从 1 开始，我们对这个序列加一，并指定到原数据中：

df.assign(row_number=df.groupby('a').cumcount()+1)
'''
   a  b  row_number
0  1  4           1
1  2  9           1
2  2  4           2
3  2  5           3
4  1  8           2
5  2  5           4
6  1  7           3
7  1  1           4
'''

接下来看看第二个思路：

(
    df.groupby('a')['a']
    .transform('rank', method='first')
    .astype(int)
)
'''
0    1
1    1
2    2
3    3
4    2
5    4
6    3
7    4
Name: a, dtype: int64
'''

分组后，选择对 a 列进行操作，用transform() 方法调用 rank()，由于每组的值（即 a）的值是相同的，我们给 rank() 参数 method='first' 就可以按先出现的顺序给出序号。最后，将这个序列转为整型，因为 rank() 返回的是浮点型的。

最后将这个序列指定到新列，也能得到最终的结果：

df.assign(row_number=df.groupby('a')['a']
          .transform('rank', method='first')
          .astype(int)
         )
# ...

以上两种方法都可以实现需求，如果只要行号第一种方法是最简单实用的。如果要按某个列的数值大小给出序号就只能用第二种方法了，比如，我们分组内按 b 的大小给出序号：

df.assign(row_number=df.groupby('a')['b']
          .transform('rank', method='first')
          .astype(int)
         )
'''
   a  b  row_number
0  1  4           2
1  2  9           4
2  2  4           1
3  2  5           2
4  1  8           4
5  2  5           3
6  1  7           3
7  1  1           1
'''

以上，就是我们案例全部的内容。

（完）

pandas 实现类似 SQL row_number() 分组行号

数据与需求

思路

代码

相关内容