Python中如何基于稀疏SKU数据构建加权无向邻接矩阵

核心问题定位 #

df.groupby(['ID']).count()输出不符合预期的原因很明确：count()统计的是分组下每列的非空值数量，你的原始表所有SKU列都填了0/1不存在空值，所以每列的计数结果就是对应ID下的交易总条数，根本不是你要的按订单折叠购买记录的效果。如果硬要靠groupby实现折叠，应该用df.groupby(['ID'])[sku_cols].sum()，但这个方法会生成稠密矩阵，数据量稍大就容易爆内存，不推荐。

针对6000个SKU、25万条交易的规模，用稀疏矩阵实现全程无冗余计算，普通笔记本就能稳定跑，不会内存崩溃，步骤如下：

实现步骤 #

1. 数据预处理 #

原始表每行只有1个SKU标记为1，先提取每条交易对应的SKU，再给订单、SKU映射连续整数索引，方便稀疏矩阵构造：

import pandas as pd
import numpy as np
from scipy.sparse import csr_matrix

# 替换成你的SKU列名列表
sku_cols = [c for c in df.columns if c.startswith('SKU')]
# 提取每条交易对应的购买SKU
df['sku'] = df[sku_cols].idxmax(axis=1)

# 生成ID到连续索引的映射
sku_ids = df['sku'].unique()
sku2idx = {s:i for i,s in enumerate(sku_ids)}
order_ids = df['ID'].unique()
order2idx = {oid:i for i,oid in enumerate(order_ids)}

n_sku = len(sku_ids)
n_order = len(order_ids)

2. 构造订单-SKU稀疏关联矩阵 #

构造二值稀疏矩阵，行对应订单，列对应SKU，值为1代表该订单购买了对应SKU：

# 稀疏矩阵坐标三元组
rows = df['ID'].map(order2idx).values
cols = df['sku'].map(sku2idx).values
vals = np.ones(len(df), dtype=np.int32)

# 生成CSR格式稀疏矩阵，内存占用仅为稠密矩阵的1%不到
order_sku = csr_matrix((vals, (rows, cols)), shape=(n_order, n_sku))

3. 矩阵乘法直接得到共现邻接矩阵 #

共现权重的计算逻辑可以直接用矩阵乘法实现：SKU共现矩阵 = 订单-SKU矩阵的转置 × 订单-SKU矩阵。乘法结果中(i,j)位置的值就是SKU_i和SKU_j共同出现的订单数，正好是你要的无向边权重：

# 稀疏矩阵乘法，全程不会展开为稠密矩阵，计算速度极快
adj = order_sku.T.dot(order_sku)
# 对角线置0，剔除SKU和自身的共现值
adj.setdiag(0)

4. （可选）转为稠密表格式输出 #

6000×6000的int32稠密矩阵仅占约144MB内存，普通设备完全可以承载，直接转成你要的表格格式即可：

adj_df = pd.DataFrame(
    adj.toarray(),
    index=sku_ids,
    columns=sku_ids
)

性能说明 #

• 无for循环遍历，核心计算靠优化过的稀疏矩阵线性代数运算，25万条数据全流程运行耗时不超过1秒
• 中间过程稀疏矩阵内存占用仅数MB，不会触发内存溢出
• 计算结果完全匹配需求：比如示例中SKU1和SKU2共同在2个订单出现，对应矩阵位置的值就是2

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