pandas 将连接字符串转为虚拟变量

看过来

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

本例我们完成一个在数据特征工程中的虚拟变量（又叫哑变量）转换问题，源数据中的特征不是独立列值，而是在字符串中以句点连接的。我们看看如何用 pandas 来解决这个问题。

数据与需求

以下是源数据：

from io import StringIO
import pandas as pd

data = '''
投票区  vote_code
A选区   S5.S17.S6
B选区       S6.S9
C选区  S21.S5.S17
'''

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

在以上数据中，vote_code 列的 S和数字为变量，它们之间用英文句点连接。需求期望将vote_code中的变量转为列，值为对应如果包含则为 1，否则为 0，即结果为：

'''
   投票区   vote_code  S17  S21  S5  S6  S9
0  A选区   S5.S17.S6    1    0   1   1   0
1  B选区       S6.S9    0    0   0   1   1
2  C选区  S21.S5.S17    1    1   1   0   0
'''

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

思路

pandas 有一个顶级方法 pd.get_dummies() 可以将分类数据转为虚拟变量，目前源数据中的分类数据都在单元格里，因此需要先将它们转为列，一个变量在一个单元格里，可以使用列表的爆炸来操作。

爆炸后使用 pd.get_dummies() 转为虚拟变量，为了保持数据结构不变，可以再用 groupby() 合并到原来的行行数。

最后，将原来的数据拼接在左边。

代码

将 vote_code 列的值通过点拆分为列表，为爆炸做准备：

(
    df.vote_code
    .str.split('.')
)
'''

0     [S5, S17, S6]
1          [S6, S9]
2    [S21, S5, S17]
Name: vote_code, dtype: object
'''

进行列表爆炸，由于爆炸的是一个 Series，所以不需要传值：

(
    df.vote_code
    .str.split('.')
    .explode()
)
'''
0     S5
0    S17
0     S6
1     S6
1     S9
2    S21
2     S5
2    S17
Name: vote_code, dtype: object
'''

可以看到，所有的列表元素都被「炸」到行上。注意一下索引标签值，保留了原值的标签，这为我们再合并提供了条件。

用 pd.get_dummies() 转为虚拟变量，由于它不是对象方法，我们用 pipe() 来调用：

(
    df.vote_code
    .str.split('.')
    .explode()
    .pipe(pd.get_dummies)
)
'''
   S17  S21  S5  S6  S9
0    0    0   1   0   0
0    1    0   0   0   0
0    0    0   0   1   0
1    0    0   0   1   0
1    0    0   0   0   1
2    0    1   0   0   0
2    0    0   1   0   0
2    1    0   0   0   0
'''

再将这个矩阵恢复为和源数据对应的行数：

(
    df.vote_code
    .str.split('.')
    .explode()
    .pipe(pd.get_dummies)
    .groupby(level=0).sum()
)
'''
   S17  S21  S5  S6  S9
0    1    0   1   1   0
1    0    0   0   1   1
2    1    1   1   0   0
'''

groupby 后按组相加即可，因为每行的列表元素不会重复。

最后将源数据拼接在左边：

(
    df.vote_code
    .str.split('.')
    .explode()
    .pipe(pd.get_dummies)
    .groupby(level=0).sum()
    .pipe(df.join)
)
'''
   投票区   vote_code  S17  S21  S5  S6  S9
0  A选区   S5.S17.S6    1    0   1   1   0
1  B选区       S6.S9    0    0   0   1   1
2  C选区  S21.S5.S17    1    1   1   0   0
'''

这样就完成了需求。

（完）

pandas 将连接字符串转为虚拟变量

数据与需求

思路

代码

相关内容