百万行Pandas DataFrame中A类两两共享B的计数优化方案问询

问题背景 #

你有一个含100万行、两列（A为学生，B为课程）的Pandas DataFrame，需要统计任意一对学生共同选修的课程数，仅保留至少共享一门课程的学生对。原方法通过生成哑变量矩阵再做矩阵乘法，存在内存占用高、计算冗余（包含对角线和无交集对）的问题。

优化思路 #

核心是利用课程分组+学生对组合统计：每个课程对应的学生列表中，所有两两无序组合即为共享该课程的学生对，统计每个组合出现的次数就是共同课程数。这种方式只处理实际有交集的学生对，完全避免稀疏矩阵的内存浪费和无效计算。

方案一：基础高效实现 #

import pandas as pd
from itertools import combinations

# 示例数据（替换为你的百万行数据即可）
df = pd.DataFrame({'A' : ['a1','a2','a3','a1','a3','a4','a2','a1','a3','a3','a4'],
 'B':['b1','b1','b2','b3','b3','b3','b4','b5','b5','b6','b6']})

# 按课程分组，获取每个课程下的唯一学生列表
course_student_groups = df.groupby('B')['A'].unique()

# 统计学生对的共同课程数
pair_counter = {}
for students in course_student_groups:
    # 生成有序的两两组合，确保(a1,a3)和(a3,a1)被视为同一对
    for pair in combinations(sorted(students), 2):
        pair_counter[pair] = pair_counter.get(pair, 0) + 1

# 转换为结构化DataFrame
result = pd.DataFrame.from_dict(
    pair_counter,
    orient='index',
    columns=['共同课程数']
)
result.index = pd.MultiIndex.from_tuples(result.index, names=['学生1', '学生2'])
print(result)

方案二：Pandas风格批量实现 #

适合超大数据量，利用Pandas内置优化减少手动循环：

import pandas as pd
from itertools import combinations

df = pd.DataFrame({'A' : ['a1','a2','a3','a1','a3','a4','a2','a1','a3','a3','a4'],
 'B':['b1','b1','b2','b3','b3','b3','b4','b5','b5','b6','b6']})

# 按课程生成所有学生对组合
course_pairs = df.groupby('B')['A'].apply(
    lambda x: pd.Series(list(combinations(sorted(x.unique()), 2)))
)

# 统计每个学生对的出现次数
result = course_pairs.value_counts().reset_index(name='共同课程数')
result.columns = ['学生1', '学生2', '共同课程数']
print(result)

方案优势 #

• 内存高效：无需生成巨大的哑变量稀疏矩阵，仅处理有共同课程的学生对
• 计算高效：跳过无交集学生对和对角线的无效计算，只统计实际存在的共享关系
• 结果精准：自动处理同一学生重复选同一课程的情况（通过unique()去重），避免重复计数

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