牛客题解 | 实现TF-IDF

题目

题目链接

TF-IDF是一种衡量文本特征的指标，常用于文本分类和信息检索。其计算公式为：

\[TF-IDF = TF \times IDF \]

其中，TF是词频，IDF是逆文档频率。
TF的计算公式为：

\[TF = \frac{词频}{文档长度} \]

IDF的计算公式为：

\[IDF = \log(\frac{文档总数}{包含该词的文档数 + 1}) + 1 \]

IDF还有很多其他的计算方式，但本题中使用的是上述公式。
在这个算法中，词汇表发挥着重要作用，是计算TF-IDF的基础。通常，词汇表是所有文档中所有单词的集合，但在本题中，词汇表是所有文档中所有单词的集合，并加上查询词。

标准代码如下

def compute_tf_idf(corpus, query):
    vocab = sorted(set(word for document in corpus for word in document).union(query))
    word_to_index = {word: idx for idx, word in enumerate(vocab)}
    tf = np.zeros((len(corpus), len(vocab)))
    for doc_idx, document in enumerate(corpus):
        for word in document:
            word_idx = word_to_index[word]
            tf[doc_idx, word_idx] += 1
        tf[doc_idx, :] /= len(document)
    df = np.count_nonzero(tf > 0, axis=0)
    num_docs = len(corpus)
    idf = np.log((num_docs + 1) / (df + 1)) + 1
    tf_idf = tf * idf
    query_indices = [word_to_index[word] for word in query]
    tf_idf_scores = tf_idf[:, query_indices]
    tf_idf_scores = np.round(tf_idf_scores, 5)
    return tf_idf_scores.tolist()

该算法是一个常见的文本处理算法，也可以使用sklearn库中的TfidfVectorizer函数来实现TF-IDF的计算