Python实现TF-IDF算法进行文本关键词提取

栗子

发布时间：2026-10-11| 3人浏览过

原创

本文介绍如何使用Python实现TF-IDF算法进行文本关键词提取。TF-IDF由词频（TF）和逆文档频率（IDF）组成，通过计算词频与逆文档频率的乘积来评估词语在文档中的重要性。文章详细讲解了TF-IDF的计算步骤，包括计算词频、逆文档频率以及最终的TF-IDF值，并提供了完整的Python代码实现。

TF-IDF算法原理

TF-IDF（Term Frequency-Inverse Document Frequency）是一种用于信息检索与数据挖掘的常用加权技术。该算法由两个核心部分组成：词频（TF）和逆文档频率（IDF）。

词频（TF）衡量一个词在当前文档中出现的频率。逆文档频率（IDF）则反映该词在整个语料库中的普遍程度。IDF值与词的常见程度成反比：一个词在越多的文档中出现，其IDF值越小。通过计算TF与IDF的乘积，可以得到TF-IDF值。该值越大，表示该词对当前文档的重要性越高。

TF-IDF算法的优点在于其简单性和高效性，且计算结果通常符合实际情况。相比之下，单纯使用词频衡量关键词不够全面，因为一些重要词汇可能在文档中出现的次数并不多。

TF-IDF计算步骤

实现TF-IDF算法主要包含以下三个步骤：

计算词频（TF）：词频等于某词在文档中出现的次数除以文档的总词数，或者除以文档中出现次数最多的词的词数。
计算逆文档频率（IDF）：计算IDF需要一个语料库来模拟语言环境。逆文档频率的计算公式为：log(文档总数 / (包含该词的文档数 + 1))。加1是为了避免分母为零的情况。
计算TF-IDF：最终的TF-IDF值等于词频乘以逆文档频率。

Python代码实现

以下Python代码实现了TF-IDF算法，包括读取文本文件、统计词频、遍历文件夹、计算TF-IDF以及将结果写入文件等功能。

#!/usr/bin/env python
#-*- coding:utf-8 -*-
 
'''
计算文档的TF-IDF
'''
import codecs
import os
import math
import shutil
 
#读取文本文件
def readtxt(path):
 with codecs.open(path,"r",encoding="utf-8") as f:
  content = f.read().strip()
 return content
 
#统计词频
def count_word(content):
 word_dic ={}
 words_list = content.split("/")
 del_word = ["\r\n","/s"," ","/n"]
 for word in words_list:
  if word not in del_word:
   if word in word_dic:
    word_dic[word] = word_dic[word]+1
   else:
    word_dic[word] = 1
 return word_dic
 
#遍历文件夹
def funfolder(path):
 filesArray = []
 for root,dirs,files in os.walk(path):
  for file in files:
   each_file = str(root+"//"+file)
   filesArray.append(each_file)
 return filesArray
 
 
#计算TF-IDF
def count_tfidf(word_dic,words_dic,files_Array):
 word_idf={}
 word_tfidf = {}
 num_files = len(files_Array)
 for word in word_dic:
  for words in words_dic:
   if word in words:
    if word in word_idf:
     word_idf[word] = word_idf[word] + 1
    else:
     word_idf[word] = 1
 for key,value in word_dic.items():
  if key !=" ":
   word_tfidf[key] = value * math.log(num_files/(word_idf[key]+1))
 
 #降序排序
 values_list = sorted(word_tfidf.items(),key = lambda item:item[1],reverse=True)
 return values_list
 
#新建文件夹
def buildfolder(path):
 if os.path.exists(path):
  shutil.rmtree(path)
 os.makedirs(path)
 print("成功创建文件夹！")
 
#写入文件
def out_file(path,content_list):
 with codecs.open(path,"a",encoding="utf-8") as f:
  for content in content_list:
   f.write(str(content[0]) + ":" + str(content[1])+"\r\n")
 print("well done!")
 
def main():
 #遍历文件夹
 folder_path = r"分词结果"
 files_array = funfolder(folder_path)
 #生成语料库
 files_dic = []
 for file_path in files_array:
  file = readtxt(file_path)
  word_dic = count_word(file)
  files_dic.append(word_dic)
 #新建文件夹
 new_folder = r"tfidf计算结果"
 buildfolder(new_folder)
 
 #计算tf-idf,并将结果存入txt
 i=0
 for file in files_dic:
  tf_idf = count_tfidf(file,files_dic,files_array)
  files_path = files_array[i].split("//")
  #print(files_path)
  outfile_name = files_path[1]
  #print(outfile_name)
  out_path = r"%s//%s_tfidf.txt"%(new_folder,outfile_name)
  out_file(out_path,tf_idf)
  i=i+1
 
if __name__ == '__main__':
 main()