pandas 按组显示列值的列表

看过来

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

大家都知道，分组对象的 transform() 方法能够返回一个与源数据同样长度的值，实现类似于 SQL 窗口函数的功能。今天的案例，让我们来进一步理解 transform() 的机制。

数据与需求

以下是源数据：

from io import StringIO
import pandas as pd

data = '''
组别  值
A   1
A   2
A   3
B   4
B   5
C   6
C   7
B   8
'''

df = pd.read_csv(StringIO(data), sep=r'\s+')
df
# ...

基于以上数据，需要增加一个名为「所属组别的值」的列，对应值为当前组别下「值」的列表，即结果为：

'''
  组别  值 所属组别的值
0  A  1  1,2,3
1  A  2  1,2,3
2  A  3  1,2,3
3  B  4  4,5,8
4  B  5  4,5,8
5  C  6    6,7
6  C  7    6,7
7  B  8  4,5,8
'''

接下来分析一下解决思路。

思路

根据以往我们的思路，要得到此列的值，先按「组别」进行分组，然后用 transform() 方法写一个函数，将「值」转为列表，但这个列表也是一个序列，它会直接填充数据，而不会广播到每个组的值中，这时我们要将它构造成一个值为列表的序列。

另外，还可以用分组后的 agg() 来聚合，再用 join() 拼接或者映射方法加到源数据后边。

代码

先分组，然后转为列表：

df.groupby('组别').值.transform(list)
'''
0    1
1    2
2    3
3    4
4    5
5    6
6    7
7    8
Name: 值, dtype: int64
'''

可以看到，聚合后的列表并没有广播，而是直接展开了，这时我们写一个方法来构造值为列表、长度为分组长度的序列：

foo = df.groupby('组别').值.transform(lambda x: [x.to_list()] * len(x))
foo
'''
0    [1, 2, 3]
1    [1, 2, 3]
2    [1, 2, 3]
3    [4, 5, 8]
4    [4, 5, 8]
5       [6, 7]
6       [6, 7]
7    [4, 5, 8]
Name: 值, dtype: object
'''

这下，我们构造的元素为列表的列表展开后填充为源数据的结构中，最后再将这列指定到源数据中：

df.assign(所属组别的值=foo)
'''
  组别  值     所属组别的值
0  A  1  [1, 2, 3]
1  A  2  [1, 2, 3]
2  A  3  [1, 2, 3]
3  B  4  [4, 5, 8]
4  B  5  [4, 5, 8]
5  C  6     [6, 7]
6  C  7     [6, 7]
7  B  8  [4, 5, 8]
'''

此时已经完成了需求。我们再来按第二种思路解决这个问题。

拼接和映射方法

先将分组后的数据聚合：

df.groupby('组别').值.agg(list)
'''
组别
A    [1, 2, 3]
B    [4, 5, 8]
C       [6, 7]
Name: 值, dtype: object
'''

得到了以「组别」为索引的数据，这时再以「组别」为键用 join() 来拼接源数据：

df.join(
    df.groupby('组别').值.agg(list).rename('所属组别的值'),
    on='组别'
)
# ... <略，见正确的结果>

为了将新增加的列有正确的名字，我们对聚合 Series 数据增加了名称属性，join() 方法的 on 参数为连接键。

还可以用 map() 来映射添加：

df.assign(所属组别的值=df.组别.map(df.groupby('组别').值.agg(list) ))
# ... <略，见正确的结果>

这是利用 Series 的字典特性来完成的操作。

这样就完成了需求。

（完）

pandas 按组显示列值的列表

数据与需求

思路

代码

拼接和映射方法

相关内容