用python编写一个TF-IDF算法的完整代码。 要求:输入输出标注清楚;输入为已分词后的txt文档,输出结果为排序后的词语及TF-IDF值,输出形式为xlsx文档;标注详细清晰;以注释形式描述所使用的公式。
时间: 2024-03-05 19:54:03 浏览: 121
人工智能-项目实践-文档相关性搜索-用python实现TF-IDF算法,用于文档的相关性搜索
好的,以下是Python编写的TF-IDF算法的完整代码:
```python
import math
import pandas as pd
from collections import Counter
# 计算TF值
def calc_tf(word_list):
word_count = Counter(word_list)
word_num = len(word_list)
tf_dict = {}
for word, count in word_count.items():
tf_dict[word] = count / word_num
return tf_dict
# 计算IDF值
def calc_idf(word_lists):
idf_dict = {}
doc_num = len(word_lists)
word_sets = set(word for word_list in word_lists for word in word_list)
for word in word_sets:
count = sum(1 for word_list in word_lists if word in word_list)
idf_dict[word] = math.log(doc_num / count)
return idf_dict
# 计算TF-IDF值
def calc_tfidf(tf_dict, idf_dict):
tfidf_dict = {}
for word, tf in tf_dict.items():
idf = idf_dict[word]
tfidf_dict[word] = tf * idf
return tfidf_dict
# 读取文本文件并分词
def read_file(file_path):
with open(file_path, 'r', encoding='utf-8') as f:
text = f.read()
word_list = text.split()
return word_list
# 将TF-IDF结果保存为Excel文件
def save_to_excel(tfidf_dict, file_path):
sorted_tfidf = sorted(tfidf_dict.items(), key=lambda x: x[1], reverse=True)
data = {'word': [], 'tf-idf': []}
for word, tfidf in sorted_tfidf:
data['word'].append(word)
data['tf-idf'].append(tfidf)
df = pd.DataFrame(data)
df.to_excel(file_path, index=None)
if __name__ == '__main__':
# 读取文本文件并分词
word_lists = []
word_lists.append(read_file('text1.txt'))
word_lists.append(read_file('text2.txt'))
# 计算TF值
tf_dicts = []
for word_list in word_lists:
tf_dict = calc_tf(word_list)
tf_dicts.append(tf_dict)
# 计算IDF值
idf_dict = calc_idf(word_lists)
# 计算TF-IDF值
tfidf_dicts = []
for tf_dict in tf_dicts:
tfidf_dict = calc_tfidf(tf_dict, idf_dict)
tfidf_dicts.append(tfidf_dict)
# 将TF-IDF结果保存为Excel文件
save_to_excel(tfidf_dicts[0], 'text1_tfidf.xlsx')
save_to_excel(tfidf_dicts[1], 'text2_tfidf.xlsx')
```
以上是完整的TF-IDF算法代码,下面是对代码中的各个函数和变量进行详细说明:
- `calc_tf(word_list)`:计算一篇文档的TF值,其中`word_list`为已分词的单词列表。
- `calc_idf(word_lists)`:计算多篇文档的IDF值,其中`word_lists`为已分词的单词列表的列表。
- `calc_tfidf(tf_dict, idf_dict)`:根据给定的TF值和IDF值计算TF-IDF值,其中`tf_dict`为一篇文档的TF值,`idf_dict`为多篇文档的IDF值。
- `read_file(file_path)`:读取指定路径下的文本文件,并分词得到单词列表。
- `save_to_excel(tfidf_dict, file_path)`:将TF-IDF值结果保存为Excel文件,其中`tfidf_dict`为计算得到的TF-IDF值字典,`file_path`为保存的文件路径。
在主函数中,我们首先读取两篇文档并分别计算它们的TF-IDF值,然后将结果保存为Excel文件。需要注意的是,这里我们使用了Pandas库来将结果保存为Excel文件,因此需要先安装好Pandas库。
阅读全文