直接用 sklearn 的 TfidfVectorizer 当然爽,但面试 / 做复杂需求时,手写一遍才能调优。

核心公式其实就两行:

tf(t, d) = 词 t 在文档 d 里的次数 / d 的总词数
idf(t) = log(文档总数 / 包含 t 的文档数 + 1)
tfidf = tf * idf

手写版:

from collections import Counter
import math

def tfidf(corpus):
    n = len(corpus)
    df = Counter()
    for doc in corpus:
        for word in set(doc):
            df[word] += 1
    result = []
    for doc in corpus:
        tf = Counter(doc)
        doc_len = len(doc)
        vec = {}
        for word, count in tf.items():
            vec[word] = (count / doc_len) * math.log(n / (df[word] + 1))
        result.append(vec)
    return result

跑一遍 sklearn 对比,结果一致。手写过一遍,后面调 max_df、min_df、ngram_range 这些参数时心里就有数了。


本文由 BBZ · 小朵科技工作室 出品。配套工具:TF-IDF 文本分析软件 P11