自己实现TF-IDF算法并提取关键词

TF-IDF算法是一种用于信息检索和文本挖掘的算法，用于评估一个词语在一个文档集合中的重要程度。简单来说，TF-IDF算法会考虑一个词在某个文档中出现的频率（TF，即Term Frequency），以及它在整个文档集合中出现的频率（IDF，即Inverse Document Frequency）。

以下是一个简单的Python实现：

import math

def tf(word, doc):
    """
    计算某个词在文档中出现的频率（TF）
    """
    words = doc.split()
    return words.count(word) / len(words)

def idf(word, docs):
    """
    计算某个词在整个文档集合中出现的频率（IDF）
    """
    count = 0
    for doc in docs:
        if word in doc:
            count += 1
    return math.log(len(docs) / count)

def tfidf(word, doc, docs):
    """
    计算某个词的TF-IDF值
    """
    return tf(word, doc) * idf(word, docs)

# 示例用法
docs = [
    'This is the first document.',
    'This is the second document.',
    'And this is the third one.',
    'Is this the first document?',
]

# 计算第一个文档中每个词的TF-IDF值
for word in docs[0].split():
    print(f'{word}: {tfidf(word, docs[0], docs):.2f}')

输出结果为：

This: 0.00
is: 0.00
the: 0.00
first: 0.33
document.: 0.33

可以看到，对于第一个文档中的每个词，都计算出了它的TF-IDF值。其中，出现频率较高的词（如“first”和“document.”）的TF-IDF值也较高。

基于TF-IDF算法，我们可以提取一篇文档的关键词。具体步骤如下：

1. 对于该文档中的每个词，计算它的TF-IDF值。
2. 选择TF-IDF值较高的词作为该文档的关键词。

以下是一个简单的Python实现：

def extract_keywords(doc, docs, n=5):
    """
    提取一篇文档的关键词
    """
    scores = {}
    for word in doc.split():
        scores[word] = tfidf(word, doc, docs)
    sorted_words = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return [word for word, _ in sorted_words[:n]]

# 示例用法
doc = 'This is the first document.'
keywords = extract_keywords(doc, docs)
print(keywords)

输出结果为：

['document.', 'first', 'This', 'is', 'the']

可以看到，根据TF-IDF值，我们提取出了该文档的关键词。