从零实现TF-IDF并按IDF值返回Top N特征的方案咨询

实现思路正确性评估 #

你的整体实现思路完全正确：先基于全量语料计算所有词汇的IDF值，按IDF降序排序后取前N个词汇构建筛选后的词汇表，符合按IDF筛选特征的核心逻辑。
你得到的前50个特征IDF完全一致属于正常现象：这些词汇的文档频率（df）都是1，也就是仅在1篇文档中出现，代入你使用的平滑IDF公式1+log((1+语料总文档数)/(1+df))，计算结果自然完全相同。

现存可优化问题 #

• IDF计算逻辑存在冗余bug：现有__IDF方法中，IDF赋值语句的缩进错误，写在了遍历文档的for循环内部，每检查一篇文档就会重新计算一次当前词的IDF，属于无效重复计算，语料量大时会严重拖慢运行速度。
• 构建Top N词汇表时使用set中转，set是无序容器，会打乱IDF排序的结果顺序，虽然不影响映射正确性，但如果需要词汇索引和IDF排序顺序对应，这个处理完全多余。
• 特征数写死为50，灵活性差，无法适配不同场景的特征数需求。
• fit方法目前返回全量词汇表和全量IDF，需要额外手动处理才能得到Top N的结果，和transform方法的调用适配性差。

更优实现方案 #

核心优化点 #

1. 参数化特征数：在类初始化时新增top_n参数，支持自定义返回的Top特征数量，默认设为None即可返回全量特征。
2. 优化IDF计算逻辑：先遍历所有文档统计每个词的文档频率，再批量计算IDF，将时间复杂度从O(词汇量*文档数)降至O(总词数)，大幅提升运行效率。
3. 简化Top N特征筛选逻辑：直接基于排序后的IDF列表取前N项构建词汇表和IDF映射，不需要set中转，保证词汇索引顺序和IDF排序顺序一致；IDF相同时增加词的字典序排序规则，保证每次运行结果稳定。
4. 对齐sklearn API规范：fit时将筛选后的词汇表和IDF值存储为实例属性，transform时直接调用实例属性，不需要额外传入参数，使用更便捷；每次调用transform时清空缓存的行列值，避免重复调用时数据累加出错。

优化后核心代码 #

from collections import Counter
from tqdm import tqdm
from scipy.sparse import csr_matrix
import math
from sklearn.preprocessing import normalize
import numpy as np

class TFIDFVectoriser_:
    # 类初始化新增top_n参数
    def __init__(self, top_n=None): 
        """这是TFIDFVectoriser类的初始化方法
        :param top_n: 按IDF降序返回的特征数量，为None时返回全量特征
        """
        self.top_n = top_n
        self.vocab = {} # 存储最终使用的词->列号映射
        self.idf_values = {} # 存储最终使用的词对应的IDF值
        self.__rows = [] 
        self.__columns =[] 
        self.__values = [] 
    
    # 统计每个词的文档频率
    def __compute_doc_freq(self, dataset):
        doc_freq = {}
        for data in dataset:
            words = set([w for w in data.split() if len(w)>=2])
            for w in words:
                doc_freq[w] = doc_freq.get(w, 0) + 1
        return doc_freq
    
    # 批量计算IDF值
    def __compute_idf(self, doc_freq, n_docs):
        idf_dict = {}
        for w, df in doc_freq.items():
            idf_dict[w] = 1 + (math.log((1 + n_docs)/(1 + df)))
        return idf_dict
    
    # 优化后的fit方法
    def fit(self,dataset):
        if not isinstance(dataset,list):
            print("输入错误：该函数仅接受列表输入，请传入词列表后重试...")
            return
        # 统计文档频率、计算IDF
        doc_freq = self.__compute_doc_freq(dataset)
        idf_values = self.__compute_idf(doc_freq, len(dataset))
        # 按IDF降序排序，IDF相同时按词升序排序保证结果稳定
        sorted_idf = sorted(idf_values.items(), key=lambda x: (-x[1], x[0]))
        # 筛选Top N特征
        if self.top_n is not None and self.top_n < len(sorted_idf):
            sorted_idf = sorted_idf[:self.top_n]
        # 构建词汇表和IDF映射
        self.vocab = {item[0]: idx for idx, item in enumerate(sorted_idf)}
        self.idf_values = {item[0]: item[1] for item in sorted_idf}
        return self.vocab, self.idf_values
    
    # 优化后的transform方法
    def transform(self,dataset):
        if not isinstance(dataset,list):
            print("输入错误：该函数仅接受列表输入，请传入词列表后重试...")
            return
        # 清空缓存的行列值，避免重复调用累加
        self.__rows.clear()
        self.__columns.clear()
        self.__values.clear()
        for idx, row in enumerate(tqdm(dataset)):
            words = row.split()
            word_freq = dict(Counter(words))
            doc_len = len(words)
            for word, freq in word_freq.items():
                if len(word) < 2:
                    continue
                col_idx = self.vocab.get(word, -1)
                if col_idx != -1:
                    self.__columns.append(col_idx)
                    self.__rows.append(idx)
                    tfidf = (freq/doc_len) * self.idf_values[word]
                    self.__values.append(tfidf)
        sparse_matrix = csr_matrix((self.__values, (self.__rows,self.__columns)), shape=(len(dataset),len(self.vocab)))
        normalized_matrix = normalize(sparse_matrix,norm='l2', axis=1, copy=True, return_norm=False)
        return normalized_matrix

使用示例 #

# 初始化时指定取Top50 IDF特征
vectorizer = TFIDFVectoriser_(top_n=50)
# 直接fit得到Top50的词汇表和IDF
vocab, idf = vectorizer.fit(你的语料列表)
# transform直接调用，不需要传额外参数
tfidf_matrix = vectorizer.transform(你的语料列表)

补充说明：如果是用于下游任务的特征选择，单纯按IDF筛选不一定是最优方案，IDF极高的词都是非常低频的词，可能携带的有效信息量有限，也可考虑结合卡方检验、互信息等特征选择方法筛选特征。

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