本文介绍如何将 DataFrame 中包含数组的列，转换为每个数组元素占据一行的新 DataFrame。通过使用 itertools.combinations 创建元素对，并结合 explode 函数，可以有效地将数组列拆分为多行，从而重塑数据结构，使其更易于分析和使用。
在数据处理中，经常会遇到 DataFrame 的某一列包含列表或数组的情况。有时，我们需要将这些列表中的每个元素拆分到单独的行中，以便进行更细粒度的分析。本文将提供一个解决方案，将 DataFrame 中数组列的元素转换为新的行，并生成元素对，解决类似“将 DataFrame 中数组元素转换为新的行”的问题。
解决方案
以下代码提供了一个通用的函数 make_pairs，用于处理包含列表的列，并将其转换为包含元素对的新列。
代码解释：
make_pairs(df: pd.DataFrame, col: str) -> pd.DataFrame 函数：
该函数接收一个 DataFrame df 和一个列名 col 作为输入。
df[col].apply(lambda x: [*combinations(iterable=x, r=2)])：对 col 列的每个列表应用 itertools.combinations 函数，生成所有可能的 2 个元素的组合。[*...] 用于将 combinations 返回的迭代器转换为列表。
.explode()：将列表中的每个元素对展开为单独的行。
pd.DataFrame(...)：使用展开后的元素对创建一个新的 DataFrame，其中 index 设置为原始 DataFrame 的索引，columns 设置为 B_m1 和 B_m2 (或 B_n1 和 B_n2)。
示例代码：
创建了一个示例 DataFrame df，其中 B_m 和 B_n 列包含列表。
df.join(...)：使用 make_pairs 函数处理 B_m 和 B_n 列，并将结果与原始 DataFrame 连接起来。
.drop(columns=["B_m", "B_n"])：删除原始的 B_m 和 B_n 列。
运行结果
运行上述代码将生成以下 DataFrame：
总结
本文提供了一种将 DataFrame 中包含数组的列转换为新行的方法。通过使用 itertools.combinations 和 explode 函数，我们可以有效地重塑数据，使其更适合分析。这种方法可以应用于各种数据处理场景，例如处理包含标签列表、特征列表或任何需要分解到单独行的数组数据的 DataFrame。
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]
# 创建每个列表中元素的 2-pair 组合
.apply(lambda x: [*combinations(iterable=x, r=2)])
# 将其展开为包含 2 个元素的列表的 Series
.explode()
)
# 使用原始索引构建 DataFrame 以进行连接
return pd.DataFrame(
data=pairs.to_list(),
index=pairs.index,
columns=[f"{col}{i}" for i in range(1, 3)]
)
# 示例 DataFrame
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'],
)
# 连接所有内容。
out = (
df.join(
other=[
make_pairs(df=df, col="B_m"),
make_pairs(df=df, col="B_n"),
],
)
# 删除不需要的列。
.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