大规模单词列表的共现矩阵生成方法选型咨询

问题背景 #

需为大规模单词列表生成单词级滑动窗口共现矩阵（统计指定窗口内两个单词共同出现的次数），现有两种实现方式，需从效率、准确性、易用性三方面评估适配性。

示例单词列表（节选）：

word_list = ['geben', 'interessieren', 'bringen', 'lassen', 'stellen', 'sehen', ..., 'gast']

原方法的核心问题 #

方法1：自定义Python代码（字符级错误实现） #

def _co_occurrence_n_gram(self, word_list, window_size):
    co_occurence_counts = defaultdict(int)
    vocab = set()

    # Iterate over the words in the word list
    for word in word_list:
        # iterate over tokens in the word —— 错误：将单个单词拆分为字符
        for i, token in enumerate(word):
            vocab.add(token)  # 收集的是字符而非单词

            # Get the next tokens within the window size
            next_tokens = word[i + 1 : i + 1 + window_size]
            for next_token in next_tokens:
                # Create a tuple of the sorted tokens and increment the co-occurrence count
                co_occurrence_key = tuple(sorted([next_token, token]))
                co_occurence_counts[co_occurrence_key] += 1

    # Create a DataFrame to represent the co-occurrence counts
    vocab = sorted(vocab)  # sort eh vocab
    data_frame = pd.DataFrame(
        data=np.zeros((len(vocab), len(vocab)), dtype=np.int16), index=vocab, columns=vocab
    )

    # populate the DataFrame with the co-occurrence counte
    for co_occurrence_key, count in co_occurence_counts.items():
        data_frame.at[co_occurrence_key[0], co_occurrence_key[1]] = count
        data_frame.at[co_occurrence_key[1], co_occurrence_key[0]] = count

    return data_frame

核心问题：代码逻辑是对单个单词的字符做共现统计，而非单词级共现，完全不符合需求。

方法2：scikit-learn CountVectorizer（全局共现错误实现） #

def _co_occurrence_n_gram(self, word_list):
    # Flatten the list of words
    all_words = " ".join(word_list)

    # Create a CountVectorizer with ngram_range=(1, 1) and stop_words='english'
    cv = CountVectorizer(ngram_range=(1, 1))

    # Convert the list of words into a matrix of token counts
    X = cv.fit_transform([all_words])

    # Calculate the co-occurrence matrix
    Xc = X.T * X
    Xc.setdiag(0)  # Set the diagonal elements to zero

    # Get the feature names (tokens)
    names = cv.get_feature_names_out()

    # Create a DataFrame to represent the co-occurrence matrix
    data_frame = pd.DataFrame(data=Xc.toarray(), columns=names, index=names)

    return data_frame

核心问题：X.T * X计算的是词频外积，而非滑动窗口内的共现。例如单词kostenlos和sender各出现2次，结果会显示共现4次，但实际仅共同出现2次，不符合滑动窗口共现定义。

修正后的正确实现对比 #

修正后的自定义方法（单词级） #

from collections import defaultdict
import pandas as pd
import numpy as np

def co_occurrence_window(word_list, window_size=2):
    co_counts = defaultdict(int)
    vocab = set(word_list)
    
    for idx, word in enumerate(word_list):
        # 确定窗口左右边界
        start = max(0, idx - window_size)
        end = min(len(word_list), idx + window_size + 1)
        # 窗口内排除当前单词的其他词汇
        window_words = word_list[start:idx] + word_list[idx+1:end]
        for neighbor in window_words:
            # 有序元组避免重复统计
            pair = tuple(sorted((word, neighbor)))
            co_counts[pair] += 1
    
    # 构建共现矩阵
    vocab_sorted = sorted(vocab)
    mat = np.zeros((len(vocab_sorted), len(vocab_sorted)), dtype=np.int32)
    word_to_idx = {word:i for i, word in enumerate(vocab_sorted)}
    
    for (w1, w2), cnt in co_counts.items():
        i = word_to_idx[w1]
        j = word_to_idx[w2]
        mat[i][j] = cnt
        mat[j][i] = cnt
    
    return pd.DataFrame(mat, index=vocab_sorted, columns=vocab_sorted)

修正后的scikit-learn+滑动窗口实现 #

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.preprocessing import MultiLabelBinarizer
import pandas as pd
import numpy as np

def co_occurrence_sklearn(word_list, window_size=2):
    # 生成所有滑动窗口的词汇集合
    windows = []
    for idx in range(len(word_list)):
        start = max(0, idx - window_size)
        end = min(len(word_list), idx + window_size + 1)
        window = word_list[start:end]
        windows.append(window)
    
    # 转换窗口为二进制矩阵
    mlb = MultiLabelBinarizer()
    window_mat = mlb.fit_transform(windows)
    # 计算共现矩阵并清空对角线
    co_mat = window_mat.T @ window_mat
    np.fill_diagonal(co_mat, 0)
    
    return pd.DataFrame(co_mat, index=mlb.classes_, columns=mlb.classes_)

多维度对比 #

1. 准确性 #

修正后的两种方法均能正确生成单词级滑动窗口共现矩阵，结果完全一致，可通过小样本测试验证。

2. 效率 #

• 自定义方法：时间复杂度为O(N*W)（N为单词列表长度，W为窗口大小），纯Python循环在百万级单词场景下速度较慢，但内存占用低，仅存储词对计数。
• scikit-learn方法：依赖numpy底层C实现的矩阵运算，速度远快于纯Python循环，但词汇量极大（十万级+）时，窗口矩阵会占用大量内存。

3. 易用性 #

• 自定义方法：逻辑直观，易调整窗口规则、过滤特定词汇，仅依赖pandas和numpy。
• scikit-learn方法：代码简洁，但滑动窗口逻辑封装较深，调整规则时需修改窗口生成部分，灵活性不如自定义方法。

推荐方案 #

• 若单词列表规模在十万级以内，优先选择scikit-learn版本，速度更快。
• 若单词列表达百万级以上或词汇量极大，优先选择自定义方法，内存占用更可控，也可通过numba对循环加速进一步提升效率。
• 无论选择哪种方案，必须先修正原代码的逻辑错误，确保实现的是单词级滑动窗口共现。

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