从词袋DataFrame转换为数组的Python高效实现方案咨询

首先得明确你的DataFrame结构——我猜它应该包含文档ID、**词汇索引（vocab_idx）和计数（count）**这几列，对吧？用groupby+apply确实会因为Python层面的循环开销拖慢速度，下面给你两种完全向量化的高效方案，比apply快几个数量级：

方案一：用稀疏矩阵构建（推荐，内存+速度双优） #

因为词袋矩阵绝大多数元素都是0，用稀疏矩阵存储不仅省内存，构建速度也远快于逐组生成数组。我们可以用scipy.sparse的coo_matrix来实现：

import pandas as pd
from scipy.sparse import coo_matrix

# 假设你的DataFrame结构如下（替换成你的实际数据）
df = pd.DataFrame({
    "doc_id": [0, 0, 1, 1, 2],
    "vocab_idx": [2, 1998, 500, 1200, 300],
    "count": [1, 2, 3, 1, 4]
})

VOCAB_LEN = 2000

# 第一步：给每个唯一的doc_id分配连续的行索引（避免doc_id不连续的情况）
unique_docs = df["doc_id"].unique()
doc_to_row = {doc: idx for idx, doc in enumerate(unique_docs)}
df["row_idx"] = df["doc_id"].map(doc_to_row)

# 第二步：用coo_matrix构建稀疏词袋矩阵
row_indices = df["row_idx"].values
col_indices = df["vocab_idx"].values
count_values = df["count"].values

# 构建稀疏矩阵，指定形状为（文档数，词汇表长度），dtype用uint8节省内存
sparse_bow = coo_matrix(
    (count_values, (row_indices, col_indices)),
    shape=(len(unique_docs), VOCAB_LEN),
    dtype="uint8"
)

# 如果确实需要稠密数组（比如某些不支持稀疏矩阵的场景），再转成numpy数组
dense_bow = sparse_bow.toarray()

为什么这方法快？ #

coo_matrix是用C底层实现的向量化操作，完全避开了Python循环的开销。而且稀疏矩阵只存储非零元素，内存占用比稠密数组小得多——如果你的文档数是10000，稠密数组要占10000*2000=20MB（uint8），但稀疏矩阵可能只占几KB。

方案二：用pandas透视表（纯pandas实现，适合小数据量） #

如果不想引入scipy依赖，用pandas的pivot_table也能实现，同样是向量化操作：

import pandas as pd

# 还是用上面的示例DataFrame
VOCAB_LEN = 2000

# 透视表：按doc_id分组，列是vocab_idx，值是count的总和，空值填0
pivot_df = df.pivot_table(
    index="doc_id",
    columns="vocab_idx",
    values="count",
    aggfunc="sum",
    fill_value=0
)

# 补全所有词汇索引列（从0到1999），确保矩阵维度是VOCAB_LEN
all_vocab = pd.RangeIndex(0, VOCAB_LEN)
pivot_df = pivot_df.reindex(columns=all_vocab, fill_value=0)

# 转成numpy数组
bow_matrix = pivot_df.values.astype("uint8")

注意事项 #

• 这个方法在词汇表很大（比如超过10000）时，透视表会占用大量内存，因为它是稠密存储的，所以更适合小数据量场景。
• 两种方案都完全避开了apply，利用了pandas/scipy的向量化优化，速度会比原来的方案快很多。

另外提醒一句：如果你的下游模型（比如sklearn的分类器、TF-IDF转换器）支持稀疏矩阵输入，直接用sparse_bow就好，不用转成稠密数组，既省内存又能加快后续计算。

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