本文旨在解决将DataFrame中包含数组的列转换为多行，每行包含数组元素组合的问题。通过使用`itertools.combinations`和`pandas`的`explode`方法，我们将展示如何有效地将DataFrame中的数组元素展开为新的行，并生成所需的组合。这对于数据重塑和分析非常有用，特别是在处理包含多个相关值的列表数据时。
在数据处理中，经常会遇到DataFrame的某一列包含列表（数组）的情况。有时，我们需要将这些列表中的元素展开，并与其他列的数据进行组合，生成新的行。本教程将介绍一种使用 itertools.combinations 结合 pandas 方法实现这一目标的方法。
问题描述
假设我们有一个DataFrame，其中包含一些列，例如 'Group', 'A_x', 'A_y'，以及两个包含列表的列 'B_m' 和 'B_n'。我们的目标是将 'B_m' 和 'B_n' 列中的列表元素进行两两组合，并将每个组合作为新的行添加到DataFrame中。
例如，对于以下DataFrame：
import pandas as pd
df = pd.DataFrame(
data=[
[0, 4, 9, [8, 7, 3], [-10, 5, 2]],
[0, 1, 2, [8, 7, 3], [-10, 5, 2]],
[1, 3, 3, [1, 2], [-5, 1]],
],
columns=['Group', 'A_x', 'A_y', 'B_m', 'B_n'],
)
print(df)输出：
Group A_x A_y B_m B_n 0 0 4 9 [8, 7, 3] [-10, 5, 2] 1 0 1 2 [8, 7, 3] [-10, 5, 2] 2 1 3 3 [1, 2] [-5, 1]
我们希望将 'B_m' 列的 [8, 7, 3] 展开为 (8, 7), (8, 3), (7, 3)，并将 'B_n' 列的 [-10, 5, 2] 展开为 (-10, 5), (-10, 2), (5, 2)，然后将这些组合添加到新的行中。
解决方案
以下代码展示了如何使用 itertools.combinations 和 pandas 方法来实现这一目标：
from itertools import combinations
import pandas as pd
def make_pairs(df: pd.DataFrame, col: str) -> pd.DataFrame:
"""
将DataFrame中指定列的列表元素进行两两组合，并返回一个新的DataFrame。
Args:
df: 原始DataFrame。
col: 需要进行组合的列名。
Returns:
一个新的DataFrame，包含组合后的元素。
"""
pairs = (
df[col]
# Create 2-pair combinations of each list in the series
.apply(lambda x: [*combinations(iterable=x, r=2)])
# Explode into a series of 2-item lists
.explode()
)
# Construct a dataframe of the using the original index for joining
return pd.DataFrame(
data=pairs.to_list(),
index=pairs.index,
columns=[f"{col}{i}" for i in range(1, 3)]
)
df = pd.DataFrame(
data=[
[0, 4, 9, [8, 7, 3], [-10, 5, 2]],
[0, 1, 2, [8, 7, 3], [-10, 5, 2]],
[1, 3, 3, [1, 2], [-5, 1]],
],
columns=['Group', 'A_x', 'A_y', 'B_m', 'B_n'],
)
# Join everything together.
out = (
df.join(
other=[
make_pairs(df=df, col="B_m"),
make_pairs(df=df, col="B_n"),
],
)
# Drop the unneeded columns.
.drop(columns=["B_m", "B_n"])
)
print(out)输出：
Group A_x A_y B_m1 B_m2 B_n1 B_n2 0 0 4 9 8 7 -10 5 0 0 4 9 8 7 -10 2 0 0 4 9 8 7 5 2 0 0 4 9 8 3 -10 5 0 0 4 9 8 3 -10 2 0 0 4 9 8 3 5 2 0 0 4 9 7 3 -10 5 0 0 4 9 7 3 -10 2 0 0 4 9 7 3 5 2 1 0 1 2 8 7 -10 5 1 0 1 2 8 7 -10 2 1 0 1 2 8 7 5 2 1 0 1 2 8 3 -10 5 1 0 1 2 8 3 -10 2 1 0 1 2 8 3 5 2 1 0 1 2 7 3 -10 5 1 0 1 2 7 3 -10 2 1 0 1 2 7 3 5 2 2 1 3 3 1 2 -5 1
代码解释
-
make_pairs 函数:
- 接收一个DataFrame和一个列名作为输入。
- 使用 df[col].apply(lambda x: [*combinations(iterable=x, r=2)]) 对指定列的每个列表应用 itertools.combinations 函数，生成所有两两组合。[*...] 用于立即解包迭代器。
- 使用 .explode() 将每个列表的组合展开为单独的行。
- 创建一个新的DataFrame，其中包含组合后的元素，并使用原始DataFrame的索引进行连接。
- 为新列添加后缀，例如 B_m1 和 B_m2。
-
主程序:
- 创建示例DataFrame。
- 调用 make_pairs 函数处理 'B_m' 和 'B_n' 列，生成包含组合的新DataFrame。
- 使用 df.join 将原始DataFrame与新DataFrame连接起来。
- 使用 .drop 删除原始的 'B_m' 和 'B_n' 列。
总结
本教程展示了如何使用 itertools.combinations 结合 pandas 方法，将DataFrame中包含数组的列转换为多行，每行包含数组元素组合。这种方法对于数据重塑和分析非常有用，特别是在处理包含多个相关值的列表数据时。通过自定义函数和灵活运用 pandas 的 apply、explode 和 join 方法，可以高效地完成复杂的数据转换任务。
注意事项
- 确保理解 itertools.combinations 函数的用法，特别是 r 参数，它决定了组合中元素的数量。
- .explode() 方法会将列表中的每个元素展开为单独的行，因此在使用前请确保了解其行为。
- 在连接DataFrame时，确保使用正确的索引，以避免数据错位。
- 根据实际需求，可以修改列名的命名方式，例如添加不同的后缀或前缀。
