如何用Pandas高效构建频率矩阵？无需自连接的解决方案

问题描述 #

给定如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame([
    ['01', 'A'],
    ['01', 'B'],
    ['01', 'C'],
    ['02', 'A'],
    ['02', 'B'],
    ['03', 'B'],
    ['03', 'C']
], columns=['id', 'category'])

需要生成如下格式的频率矩阵（矩阵中(i,j)的值表示同时包含类别i和j的id数量）：

A   B   C           
A   2   2   1
B   2   3   2
C   1   2   2

当前通过自连接+透视表的方法可以实现，但会大幅增加数据量：

result = df.merge(df, on='id')
pd.pivot_table(
    result,
    index='category_x',
    columns='category_y',
    values='id',
    aggfunc='count'
)

需要一种无需自连接的高效实现方式，且直接使用pivot_table(df, index='category', columns='category', ...)无法生效。

高效实现方法 #

方法1：One-Hot编码 + 矩阵乘法 #

先对category列做one-hot编码，按id分组求和得到每个id对应的类别存在矩阵，再将该矩阵与其转置相乘，即可得到频率矩阵：

# 按id分组，对category做one-hot编码并求和
one_hot = pd.get_dummies(df['category']).groupby(df['id']).sum()
# 矩阵乘法计算共现频率
freq_matrix = one_hot.T.dot(one_hot)
print(freq_matrix)

输出结果：

A  B  C
category        
A        2  2  1
B        2  3  2
C        1  2  2

方法2：交叉表 + 矩阵乘法 #

先构建每个id和category的存在标记，再通过交叉表得到id-category的二维表，再做矩阵转置相乘：

# 生成id和category的交叉表，值为1表示该id包含该category
cross = pd.crosstab(df['id'], df['category'])
# 转置后相乘得到频率矩阵
freq_matrix = cross.T.dot(cross)
print(freq_matrix)

这两种方法都避免了自连接带来的数据膨胀，计算效率更高，尤其是在数据量较大时优势明显，且结果完全符合需求。

内容的提问来源于stack exchange，提问作者wong.lok.yin