大型Pandas DataFrame唯一三元组统计优化：PyTorch加速咨询

针对你遇到的O(n³)复杂度瓶颈，利用元素上限固定为2000的特点，结合PyTorch的批量张量操作可以大幅提升计算速度，核心思路是用GPU并行替代Python嵌套循环，以下是具体实现方案：

核心优化逻辑 #

你的需求是统计所有i<j<k位置对应的三元组，且将(bi,bj,bk)与(bj,bi,bk)视为同一三元组（即前两个元素排序后统一表示）。直接遍历每个三元组的Python循环效率极低，而PyTorch的底层CUDA操作可以批量处理这些计算，将耗时从Python级别的循环降到硬件加速的张量操作级别。

具体实现代码 #

import torch
import pandas as pd

# 假设你的DataFrame已加载为A
# 1. 将history列转换为PyTorch张量，移至GPU（如果可用）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
history_tensors = [torch.tensor(hist, dtype=torch.long).to(device) for hist in A['history']]

# 2. 批量处理所有样本的三元组
all_normalized_triples = []
for hist in history_tensors:
    seq_len = hist.size(0)
    if seq_len < 3:
        continue  # 长度不足3的样本跳过
    # 生成所有i<j<k的索引组合
    indices = torch.combinations(torch.arange(seq_len), r=3)
    # 获取对应的三元组元素
    raw_triples = hist[indices]  # shape: (C(seq_len,3), 3)
    # 对每个三元组的前两个元素排序，统一表示
    raw_triples[:, :2] = torch.sort(raw_triples[:, :2], dim=1)[0]
    all_normalized_triples.append(raw_triples)

# 3. 合并所有三元组并统计出现次数
triple_counts = {}
if all_normalized_triples:
    all_triples = torch.cat(all_normalized_triples)
    # 获取唯一三元组及对应次数
    unique_triples, counts = torch.unique(all_triples, dim=0, return_counts=True)
    # 转换为字典格式
    triple_counts = {
        (a.item(), b.item(), c.item()): cnt.item()
        for a, b, c, cnt in zip(unique_triples[:,0], unique_triples[:,1], unique_triples[:,2], counts)
    }

# 输出结果示例
print("Top 5三元组统计：")
for triple, cnt in sorted(triple_counts.items(), key=lambda x: x[1], reverse=True)[:5]:
    print(f"{triple}: {cnt}次")

关键优化点说明 #

1. 批量生成三元组：torch.combinations直接生成所有符合i<j<k的索引组合，比Python循环生成快数倍，且支持GPU加速。
2. 统一三元组表示：通过torch.sort对每个三元组的前两个元素排序，自动将(bi,bj,bk)与(bj,bi,bk)合并为同一键。
3. 高效统计：torch.unique(dim=0)批量去重并计数，避免了Python字典的逐元素累加操作，底层用CUDA实现，效率极高。

性能对比 #

• 原Python循环实现：复杂度为O(total_triples)，但每个循环步骤都是Python级别的慢操作。
• PyTorch实现：同样是O(total_triples)复杂度，但所有核心操作都是硬件加速的张量操作，在GPU上可获得10~100倍的速度提升（取决于GPU性能和数据规模）。

内容的提问来源于stack exchange，提问作者Adam