本文探讨了如何在pandas DataFrame中，针对每个分组内的每一行数据，高效地将其与同组内所有其他行的数据进行交叉组合并扩展为新的列。通过结合Pandas的groupby().apply()和numpy的数组滚动索引技术，我们能够以高性能的方式实现这种复杂的数据转换，避免了低效的循环和合并操作，适用于需要生成组内两两比较或交互特征的场景。
挑战：分组内行数据的全交叉组合
在数据分析和特征工程中，我们经常会遇到这样的需求：给定一个按某个键（例如raceid）分组的数据集，对于组内的每一条记录，我们希望能够将同组内所有其他记录的特定信息作为新的列添加到当前记录中。例如，在一个赛马数据集中，我们可能希望为每匹马的记录添加同场比赛中所有其他马匹的排名、体重等信息，以便进行更深入的分析或构建复杂的特征。
考虑以下原始数据结构，它代表了一场赛马中的六匹马：
import pandas as pd import numpy as np data_orig = { 'meetingId': [178515] * 6, 'raceId': [879507] * 6, 'horseId': [90001, 90002, 90003, 90004, 90005, 90006], 'position': [1, 2, 3, 4, 5, 6], 'weight': [51, 52, 53, 54, 55, 56], } data_orig_df = pd.DataFrame(data_orig) print("原始数据：") print(data_orig_df)
期望的输出是这样的：对于第一行（horseId 90001），它将包含所有六匹马的信息，其中它自己的信息作为 _1 后缀的列，第二匹马的信息作为 _2 后缀的列，依此类推。对于第二行（horseId 90002），它自己的信息作为 _1 后缀的列，而其他马匹的信息则相应地滚动填充。
# 期望输出的简化示例结构（部分列） # horseId_1 position_1 weight_1 horseId_2 position_2 weight_2 ... horseId_6 position_6 weight_6 # 90001 1 51 90002 2 52 ... 90006 6 56 # 90002 2 52 90003 3 53 ... 90001 1 51 # ...
直接使用循环和pd.merge虽然能够实现，但在处理大型数据集和多个分组时，其性能会非常低下。
核心解决方案：利用NumPy的滚动索引
为了高效地实现这种分组内的行数据全交叉组合，我们可以结合Pandas的groupby().apply()方法和NumPy强大的数组索引能力。关键在于创建一个能够“滚动”或“循环移位”数组内容的索引机制。
1. 定义滚动函数
首先，我们定义一个名为roll的函数，它接收一个DataFrame组（不包含分组键），并对其进行操作。
def roll(g): """ 对DataFrame组内的数值进行滚动索引，实现行数据的全交叉组合。 参数: g (pd.DataFrame): 组内数据，不包含分组键。 返回: pd.DataFrame: 经过滚动和扩展后的DataFrame。 """ # 将DataFrame转换为NumPy数组，便于高效操作 a = g.to_numpy() num_rows = len(a) # 创建一个索引数组，用于生成滚动效果 # x = [0, 1, 2, ..., num_rows-1] x = np.arange(num_rows) # 核心：生成滚动索引 # (x[:,None] + x) 创建一个 num_rows x num_rows 的矩阵， # 每一行表示相对于原始行的偏移量。 # 例如，对于 num_rows=6: # [[0, 1, 2, 3, 4, 5], # [1, 2, 3, 4, 5, 6], # [2, 3, 4, 5, 6, 7], # [3, 4, 5, 6, 7, 8], # [4, 5, 6, 7, 8, 9], # [5, 6, 7, 8, 9, 10]] # # % num_rows 实现循环（滚动）效果 # 例如，对于 num_rows=6: # [[0, 1, 2, 3, 4, 5], # [1, 2, 3, 4, 5, 0], # [2, 3, 4, 5, 0, 1], # [3, 4, 5, 0, 1, 2], # [4, 5, 0, 1, 2, 3], # [5, 0, 1, 2, 3, 4]] # # .ravel() 将这个二维索引矩阵展平为一维数组，用于对原始数组 `a` 进行索引。 # 例如，展平后为 [0,1,2,3,4,5, 1,2,3,4,5,0, ...] # # a[...] 使用展平的索引从原始数组 `a` 中提取数据。 # 例如，a[0], a[1], ..., a[5], a[1], a[2], ..., a[0], ... # # .reshape(num_rows, -1) 将结果重新塑形。 # num_rows 保持原始行数，-1 表示列数自动计算，它会是原始列数 * num_rows。 rolled_data = a[((x[:,None] + x) % num_rows).ravel()].reshape(num_rows, -1) # 生成新的列名 # 例如，如果原始列是 ['horseId', 'position', 'weight'] # 那么新列名将是 ['horseId_1', 'position_1', 'weight_1', # 'horseId_2', 'position_2', 'weight_2', ...] new_columns = [f'{col}_{i+1}' for i in x for col in g.columns] # 将NumPy数组转换回DataFrame，并保留原始索引 return pd.DataFrame(rolled_data, index=g.index, columns=new_columns)
2. 应用 groupby().apply()
有了 roll 函数，我们就可以将其应用到分组后的DataFrame上。
# 定义分组键 group_cols = ['meetingId', 'raceId'] # 执行分组、应用滚动函数并重置索引 output_df = (data_orig_df.groupby(group_cols) .apply(lambda g: roll(g.drop(columns=group_cols))) # 对每个组应用roll函数，注意要先移除分组键 .reset_index(group_cols) # 将分组键重新添加为普通列 ) print("n处理后的数据：") print(output_df)
结果展示
运行上述代码，将得到以下输出（与期望的 data_new 结构一致，只是列名后缀从字母变为数字，这更具通用性）：
处理后的数据： meetingId raceId horseId_1 position_1 weight_1 horseId_2 position_2 weight_2 horseId_3 position_3 weight_3 horseId_4 position_4 weight_4 horseId_5 position_5 weight_5 horseId_6 position_6 weight_6 0 178515 879507 90001 1 51 90002 2 52 90003 3 53 90004 4 54 90005 5 55 90006 6 56 1 178515 879507 90002 2 52 90003 3 53 90004 4 54 90005 5 55 90006 6 56 90001 1 51 2 178515 879507 90003 3 53 90004 4 54 90005 5 55 90006 6 56 90001 1 51 90002 2 52 3 178515 879507 90004 4 54 90005 5 55 90006 6 56 90001 1 51 90002 2 52 90003 3 53 4 178515 879507 90005 5 55 90006 6 56 90001 1 51 90002 2 52 90003 3 53 90004 4 54 5 178515 879507 90006 6 56 90001 1 51 90002 2 52 90003 3 53 90004 4 54 90005 5 55
注意事项与优化
- 性能优势：此方法利用NumPy的矢量化操作，避免了python层面的显式循环，因此在处理大规模数据集时，其性能远超基于iterrows()和pd.merge()的方案。
- 内存消耗：这种数据扩展方式会显著增加DataFrame的列数。如果原始组内元素数量较多，生成的DataFrame会非常宽，可能导致巨大的内存消耗。在实际应用中，需要根据具体需求和系统资源评估其可行性。
- 列名约定：生成的列名（如horseId_1, position_2）清晰地表明了数据来源。_1通常表示该行自身的数据，_2表示滚动一位后的数据，以此类推。可以根据实际需求调整roll函数中的列名生成逻辑。
- 适用场景：
- 特征工程：创建复杂的交互特征，例如，预测一匹马的表现时，同时考虑同场竞技的其他马匹的属性。
- 组内比较：在组内进行两两比较分析。
- 数据重塑：将组内数据从长格式转换为宽格式，但不仅仅是简单的透视，而是带有特定顺序和组合的扩展。
- 分组键处理：在apply函数内部，我们通过g.drop(columns=group_cols)将分组键从要进行滚动操作的数据中移除，以避免对这些固定值进行不必要的滚动。reset_index(group_cols)则确保最终结果中保留了这些分组信息。
总结
通过巧妙地结合Pandas的groupby().apply()和NumPy的数组滚动索引技术，我们可以高效且优雅地解决分组内行数据全交叉组合的问题。这种方法不仅提供了强大的数据转换能力，也充分利用了底层库的性能优势，是处理复杂数据重塑和特征工程任务的有效策略。然而，在应用时务必关注其潜在的内存消耗，并根据具体业务需求调整。
评论（已关闭）
评论已关闭