如何用Python基于含userID/artistID/weight的.dat文件构建稀疏矩阵

你遇到的问题很典型：用Python循环处理10万行数据速度太慢，而稀疏矩阵的索引又用错了地方。下面给你两种针对性的高效方案，分别对应不同的场景：

方案一：用Pandas透视表生成带NaN的矩阵（直观高效） #

如果你的userID和artistID的取值范围不是特别大（比如最大ID不超过几万），直接用Pandas的pivot_table是最省心的方法——它内部是向量化操作，比Python循环快几个数量级，生成的结果直接就是你要的、缺失值为NaN的矩阵：

import pandas as pd
import numpy as np

# 读取数据（你原来的读取方式没问题）
data = pd.read_table('train.dat')

# 生成透视表：行=userID，列=artistID，值=weight，缺失值填充为NaN
user_artist_matrix = data.pivot_table(
    index='userID',
    columns='artistID',
    values='weight',
    fill_value=np.nan  # 显式指定缺失值为NaN，默认也是NaN
)

# 访问某个用户-艺术家的权重，比如user45和artist7
print(user_artist_matrix.loc[45, 7])

这个方法的优势是直观，生成的DataFrame可以直接用loc按原始ID索引，完全符合你的需求。

方案二：用SciPy稀疏矩阵（适合大维度数据） #

如果你的userID或artistID的取值跨度极大（比如最大userID是100万），生成全量矩阵会浪费大量内存，这时候稀疏矩阵是更好的选择。你之前的错误在于用了DataFrame的默认行索引，而非原始的userID/artistID，正确的做法是先把原始ID映射为连续的整数索引（稀疏矩阵要求行/列是连续整数）：

import pandas as pd
import scipy.sparse as sps
import numpy as np

data = pd.read_table('train.dat')

# 1. 把原始userID和artistID映射为连续的整数索引（避免稀疏矩阵出现大量空行/列）
unique_users = data['userID'].unique()
unique_artists = data['artistID'].unique()

user_to_idx = {uid: i for i, uid in enumerate(unique_users)}
artist_to_idx = {aid: j for j, aid in enumerate(unique_artists)}

# 2. 转换为稀疏矩阵需要的行、列、值数组
rows = data['userID'].map(user_to_idx).values
cols = data['artistID'].map(artist_to_idx).values
weights = data['weight'].values

# 3. 创建COO稀疏矩阵并转为CSR格式（方便快速索引）
sparse_matrix = sps.coo_matrix(
    (weights, (rows, cols)),
    shape=(len(unique_users), len(unique_artists))
).tocsr()

# 4. 封装一个函数，按原始ID获取权重（缺失返回NaN）
def get_weight(user_id, artist_id):
    if user_id not in user_to_idx or artist_id not in artist_to_idx:
        return np.nan
    row_idx = user_to_idx[user_id]
    col_idx = artist_to_idx[artist_id]
    val = sparse_matrix[row_idx, col_idx]
    return val if val != 0 else np.nan

# 测试：获取user45和artist7的权重
print(get_weight(45, 7))

为什么你之前的方法行不通？ #

• numpy循环填充：Python的for循环遍历10万行效率极低，而Pandas/SciPy的操作都是底层C实现的，速度快得多。
• 稀疏矩阵索引错误：你之前用data.index.labels是取DataFrame的默认行索引（0,1,2...），而非原始的userID/artistID，导致矩阵的行/列和实际ID完全不对应，自然会出现索引问题。

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