pandas 组合列值并扩展为多行

看过来

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

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
    # 两个 Series 长度不同是因为标签对齐机制
    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 组合列值并扩展为多行

数据与需求

思路

代码

相关内容