能否向量化嵌套循环代码以优化共现矩阵更新？求提速建议

你的三重循环实现逻辑清晰，但处理大规模语料时效率极低——Python循环本身开销大，加上每次循环都要做字典查找和矩阵元素的零散更新，会严重拖慢速度。以下是基于NumPy向量化操作的优化方案，能大幅提升运行效率：

核心优化思路 #

1. 提前完成词到索引的转换：避免在循环中重复查询字典，一次性把所有文档转换成索引数组。
2. 用数组操作替代嵌套循环：利用NumPy的广播、数组拼接等操作生成所有符合窗口条件的词对索引，再批量更新矩阵。
3. 批量累加处理重复索引：使用np.add.at处理重复的(row, col)索引，确保每个共现对的计数正确累加，这比逐个元素更新高效得多。

优化后的代码 #

步骤1：预处理语料，转换为索引序列 #

import numpy as np

# 提前将所有文档转换为词索引数组，避免循环中重复查字典
word2Ind = {words[i]: i for i in range(len(words))}
corpus_indices = [np.array([word2Ind[word] for word in doc]) for doc in corpus]

步骤2：向量化生成共现对并更新矩阵 #

M = np.zeros((num_words, num_words), dtype=np.int64)
window_size = 2  # 替换为你的实际窗口大小

# 遍历每个文档的索引数组
for doc_idx in corpus_indices:
    doc_len = len(doc_idx)
    # 计算每个位置i对应的窗口起始、结束索引
    window_starts = np.maximum(0, np.arange(doc_len) - window_size)
    window_ends = np.minimum(doc_len, np.arange(doc_len) + window_size + 1)
    
    # 生成每个i对应的所有j索引，以及对应的i索引（重复对应次数）
    j_indices = np.concatenate([np.arange(s, e) for s, e in zip(window_starts, window_ends)])
    i_indices = np.repeat(np.arange(doc_len), window_ends - window_starts)
    
    # 过滤掉i == j的情况（排除自身共现）
    valid_mask = i_indices != j_indices
    target_i = doc_idx[i_indices[valid_mask]]
    target_j = doc_idx[j_indices[valid_mask]]
    
    # 批量更新共现矩阵
    np.add.at(M, (target_i, target_j), 1)

进一步优化：合并所有文档的索引对，一次性更新 #

如果语料中有大量短文档，可以把所有文档的共现对索引合并后一次性处理，减少Python循环的开销：

all_target_i = []
all_target_j = []

for doc_idx in corpus_indices:
    doc_len = len(doc_idx)
    window_starts = np.maximum(0, np.arange(doc_len) - window_size)
    window_ends = np.minimum(doc_len, np.arange(doc_len) + window_size + 1)
    
    j_indices = np.concatenate([np.arange(s, e) for s, e in zip(window_starts, window_ends)])
    i_indices = np.repeat(np.arange(doc_len), window_ends - window_starts)
    
    valid_mask = i_indices != j_indices
    all_target_i.append(doc_idx[i_indices[valid_mask]])
    all_target_j.append(doc_idx[j_indices[valid_mask]])

# 合并所有索引对并一次性更新矩阵
all_target_i = np.concatenate(all_target_i)
all_target_j = np.concatenate(all_target_j)
np.add.at(M, (all_target_i, all_target_j), 1)

注意事项 #

• 确保word2Ind包含语料中所有词，且num_words等于词汇表的大小，避免索引越界。
• 如果语料规模极大，合并后的索引数组可能占用过多内存，此时可以分批次处理（比如每处理1000个文档就更新一次矩阵，然后清空临时列表）。

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