pandas 组合列值并扩展为多行

时间：2023-04-16 13:09:42 标签：pandas python 组合 爆炸

看过来

《pandas 教程》持续更新中，可作为 pandas 入门进阶课程、pandas 中文手册、用法大全，配有案例讲解和速查手册。提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。查看更新日志。作者开办 Python 数据分析培训，详情 Python 数据分析培训。

本教程作者所著新书《Python之光：Python编程入门与实战》（ISBN：9787111729891）已由机械工业出版社出版上市，各大电商平台有售，欢迎：查看详情并关注购买。

本教程作者所著新书《深入浅出Pandas：利用Python进行数据处理与分析》（ISBN：9787111685456）已由机械工业出版社出版上市，各大电商平台有售，欢迎：查看详情并关注购买。

pandas 处理数据时会遇到按一定的规则扩展的需求，在本例中，我们将先对数据按情况进行组合，然后展开为多行数据，这其中会涉及复杂的数据组合。

数据与需求

我们构造一个简单的数据如下：

import pandas as pd

data = {'id': ['i1', 'i2', 'i3'],
        'A': ['a1', 'a2', 'a3'],
        'C': ['', '', 'c3',],
        'D': ['d1', 'd2', 'd3'],
        'E': ['', 'e2', 'e3']}

df = pd.DataFrame(data)
df
'''
   id   A   C   D   E
0  i1  a1      d1    
1  i2  a2      d2  e2
2  i3  a3  c3  d3  e3
'''

需求为对以上数据进行扩展，id 列内容不变，增加 foo 列，foo 的内容由 A 列开头与后续其他列组合而成，不同列值之间用逗号隔开。结果如下：

'''
   id          foo
0  i1        a1,d1
1  i2        a2,d2
1  i2        a2,e2
1  i2     a2,d2,e2
2  i3        a3,c3
2  i3        a3,d3
2  i3        a3,e3
2  i3     a3,c3,d3
2  i3     a3,c3,e3
2  i3     a3,d3,e3
2  i3  a3,c3,d3,e3
'''

思路

总体思路为利用 apply 按行调用每个由行形成的 Series，然后生成一个组合数据的列表，最后将这个列表用 explode 方法爆炸最终形成需要的结果数据。

关键的是如何编写 apply 中和处理函数，它传入一个每行形成的 Series，返回组合后的列表。先排除 C 列和为空字符串的列，形成新的可用 Series，然后用这个 Series 进行组合。利用 Python 标准库中 itertools.combinations 方便来完成组合。

最后将组合中的每个排列增加 A 列值为开头，并用逗号拼接起来返回。

代码

根据思路，我们编写的处理函数如下：

import itertools

def func(ser: pd.Series):
    head = ser.A
    part = ser.loc['C':].loc[~(ser == '')]
    res = []

    for n in range(1, len(part)+1):
        it = itertools.combinations(part, n)
        res += it

    s = [','.join([ser.A]+[*i]) for i in res]

    return s

按行应用这个函数，将结果增加到 id 列的后边：

df[['id']].assign(foo=df.apply(func, axis=1))
'''
   id                                                foo
0  i1                                            [a1,d1]
1  i2                           [a2,d2, a2,e2, a2,d2,e2]
2  i3  [a3,c3, a3,d3, a3,e3, a3,c3,d3, a3,c3,e3, a3,d...
'''

爆炸 foo 列：

(
    df[['id']].assign(foo=df.apply(func, axis=1))
    .explode('foo')
)
'''
   id          foo
0  i1        a1,d1
1  i2        a2,d2
1  i2        a2,e2
1  i2     a2,d2,e2
2  i3        a3,c3
2  i3        a3,d3
2  i3        a3,e3
2  i3     a3,c3,d3
2  i3     a3,c3,e3
2  i3     a3,d3,e3
2  i3  a3,c3,d3,e3
'''

这样就完成了需求。

（完）

pandas 组合列值并扩展为多行

数据与需求

思路

代码

相关内容