如何高效实现Pandas DataFrame自组合、新值计算与排序？

原始DataFrame结构 #

+------------+-------------------+---------+--------+----------+
| process_no | process_durations | columns | orders | customer |
+------------+-------------------+---------+--------+----------+
| 0          | 3                 | [0]     | [3109] | [0]      |
+------------+-------------------+---------+--------+----------+
| 1          | 100               | [11]    | [5855] | [0]      |
+------------+-------------------+---------+--------+----------+
| 2          | 81                | [8]     | [5304] | [0]      |
+------------+-------------------+---------+--------+----------+

需求说明 #

• 将DataFrame与自身进行无重复两两组合（即combinations而非permutations）
• 计算三个衍生字段：
  • order_length：两个组合行的orders列表总长度的负值
  • column_length：两个组合行的columns列表合并后的唯一值数量
  • customers：两个组合行的customer列表合并后的唯一值列表
• 按order_length降序、column_length升序排序
• 保留原始组合的完整行数据

当前实现代码 #

from itertools import combinations
from operator import itemgetter
import pandas as pd

data = pd.DataFrame({
     'process_no': [0, 1, 2],
     'process_durations': [3, 100, 81],
     'columns': [[0], [11], [8]],
     'orders': [[3109], [5855], [5304]],
     'customer': [[0], [0], [0]]
})

vals = data.values.tolist()
cross_combine = list(combinations(vals, r=2))

sorted_cross_combine = sorted(
    [
        (
            x,
            -(len(x[0][3]) + len(x[1][3])),
            len(set(x[0][2] + x[1][2])),
            list(set(x[0][4] + x[1][4]))
        )
        for x in cross_combine
    ],
    key=itemgetter(1, 2)
)
print(sorted_cross_combine)

当前输出 #

[(([0, 3, [0], [3109], [0]], [1, 100, [11], [5855], [0]]), -2, 2, [0]),
 (([0, 3, [0], [3109], [0]], [2, 81, [8], [5304], [0]]), -2, 2, [0]),
 (([1, 100, [11], [5855], [0]], [2, 81, [8], [5304], [0]]), -2, 2, [0])]

预期输出格式 #

+-----------------------------+-----------------------------+-------------+--------------+-----------+
| x1                          | x2                          |order_length |column_length | customers |
+-----------------------------+-----------------------------+-------------+--------------+-----------+
| [0, 3, [0], [3109], [0]]    | [1, 100, [11], [5855], [0]] | -2          | 2            | [0]       |
+-----------------------------+-----------------------------+-------------+--------------+-----------+
| [0, 3, [0], [3109], [0]]    | [2, 81, [8], [5304], [0]]   | -2          | 2            | [0]       |
+-----------------------------+-----------------------------+-------------+--------------+-----------+
| [1, 100, [11], [5855], [0]] | [2, 81, [8], [5304], [0]]   | -2          | 2            | [0]       |
+-----------------------------+-----------------------------+-------------+--------------+-----------+

计算规则 #

• order_length = -(len(行1['orders']) + len(行2['orders']))
• column_length = len(set(行1['columns'] + 行2['columns']))
• customers = list(set(行1['customer'] + 行2['customer']))

性能问题 #

当数据集包含约6500条记录时，当前实现生成排序后的结果需耗时约20秒，主要瓶颈在于：

1. 重复计算每行的orders长度、columns和customer的集合
2. 列表相加后转集合的操作效率较低

高效解决方案 #

核心优化思路：预计算基础统计值 #

提前计算每行的orders长度、columns集合、customer集合，避免组合时重复计算，同时利用集合的union操作替代列表相加转集合，提升效率。

优化后代码 #

import pandas as pd
from itertools import combinations
from operator import itemgetter
import ast

# 读取并预处理数据
data = pd.read_csv('a.csv')
for col in ['columns', 'orders', 'customer']:
    data[col] = data[col].apply(ast.literal_eval)

# 预计算每行的基础统计值，避免重复计算
data['order_len'] = data['orders'].apply(len)
data['columns_set'] = data['columns'].apply(set)
data['customer_set'] = data['customer'].apply(set)

# 转换为包含原始行数据+预计算值的元组列表
vals = data.apply(
    lambda row: (row.values.tolist(), row['order_len'], row['columns_set'], row['customer_set']),
    axis=1
).tolist()

# 组合+计算：用生成器表达式减少内存占用
sorted_result = sorted(
    (
        (
            x[0],  # 行1原始数据
            y[0],  # 行2原始数据
            -(x[1] + y[1]),  # order_length
            len(x[2].union(y[2])),  # column_length
            list(x[3].union(y[3]))  # customers
        )
        for x, y in combinations(vals, r=2)
    ),
    key=itemgetter(2, 3)  # 按order_length降序、column_length升序排序
)

# 转换为预期格式的DataFrame
output_df = pd.DataFrame(
    sorted_result,
    columns=['x1', 'x2', 'order_length', 'column_length', 'customers']
)

print(output_df)

优化点说明 #

1. 预计算：只对每行计算一次order_len、columns_set、customer_set，组合时直接调用，避免重复计算
2. 集合Union：用set.union()替代list+list再转set，时间复杂度更低（集合Union操作是针对哈希表的直接合并）
3. 生成器替代列表：直接将生成器传给sorted，无需提前存储所有组合结果，大幅降低内存占用

额外性能建议 #

如果数据集持续增大，可考虑：

• 用numba对核心计算逻辑进行JIT编译
• 分批次处理组合结果，避免内存溢出

内容的提问来源于stack exchange，提问作者E. Zeytinci