如何加速25万行稀疏分类数据的自定义序数编码

问题背景 #

我通过pd.read_csv加载了CSV数据，结构如下：

其中A、B列均为无重复的分类变量列表，需要转换成如下指示表格式：

转换规则：变量出现在A列时编码为1，出现在B列时编码为-1，两列均未出现则为0。当前数据规模为25万行、25000个唯一类别，内存是主要瓶颈。我编写的代码运行耗时过长：

sparse_A = pd.get_dummies(data["A"].explode(), sparse=True).groupby(level=0).sum()
sparse_B = pd.get_dummies(data["B"].explode(), sparse=True).groupby(level=0).sum()*-1
indicator_df = sparse_A + sparse_B

求更高效的实现方式？

优化方案 #

核心思路：跳过冗余中间步骤，直接构建稀疏矩阵 #

原代码的explode+get_dummies+groupby流程会生成大量中间稀疏矩阵，分组求和更是大规模数据下的性能瓶颈。以下是两种更高效的实现方式：

方式1：向量化生成稀疏矩阵（推荐） #

完全用向量化操作替代循环，避免iterrows的开销：

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

# 1. 提取所有唯一类别并建立映射
all_cats = pd.concat([data['A'].explode(), data['B'].explode()]).unique()
cat_to_idx = {cat: idx for idx, cat in enumerate(all_cats)}
n_rows = len(data)
n_cats = len(all_cats)

# 2. 处理A列：生成行索引、列索引、值数组
a_exploded = data['A'].explode()
a_rows = a_exploded.index.values
a_cols = a_exploded.map(cat_to_idx).values
a_vals = np.ones(len(a_exploded), dtype=np.int8)  # 用int8节省内存

# 3. 处理B列：生成行索引、列索引、值数组
b_exploded = data['B'].explode()
b_rows = b_exploded.index.values
b_cols = b_exploded.map(cat_to_idx).values
b_vals = np.full(len(b_exploded), -1, dtype=np.int8)

# 4. 合并数组并构建CSR稀疏矩阵
rows = np.concatenate([a_rows, b_rows])
cols = np.concatenate([a_cols, b_cols])
vals = np.concatenate([a_vals, b_vals])

sparse_indicator = csr_matrix((vals, (rows, cols)), shape=(n_rows, n_cats))

# 可选：转换为稀疏DataFrame（如需保留DataFrame格式）
indicator_df = pd.DataFrame.sparse.from_spmatrix(
    sparse_indicator,
    index=data['idx'],
    columns=all_cats
)

方式2：循环构建（适合小批量调试） #

如果需要逐行处理的灵活性，可以用循环，但注意仅在向量化方案不适用时使用：

import pandas as pd
from scipy.sparse import csr_matrix

all_cats = pd.concat([data['A'].explode(), data['B'].explode()]).unique()
cat_to_idx = {cat: idx for idx, cat in enumerate(all_cats)}
n_rows = len(data)
n_cats = len(all_cats)

rows = []
cols = []
vals = []

for row_idx, row in data.iterrows():
    # 处理A列的类别
    for cat in row['A']:
        rows.append(row_idx)
        cols.append(cat_to_idx[cat])
        vals.append(1)
    # 处理B列的类别
    for cat in row['B']:
        rows.append(row_idx)
        cols.append(cat_to_idx[cat])
        vals.append(-1)

sparse_indicator = csr_matrix((vals, (rows, cols)), shape=(n_rows, n_cats))

优化点说明 #

• 内存优化：使用np.int8存储值（仅需1字节/元素），比默认的int64节省8倍内存；用CSR稀疏矩阵存储，仅保留非零值，避免全零列的内存浪费。
• 性能提升：跳过两次groupby.sum()操作，这一步在25万行数据下会产生大量计算开销；向量化操作比循环快一个数量级以上。
• 避免冗余操作：直接合并A、B列的编码逻辑，不需要分别生成两个稀疏矩阵再相加。

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