TF-IDF算法-Python实现（附源代码）

From: http://blog.csdn.net/zhb_bupt/article/details/40985831

一、背景

TF-IDF算法全称 termfrequency–inverse document frequency，是一种用于资讯检索与资讯探勘的常用加权技术。它的算法复杂度并不高，但能很好的满足搜索高相关度文档的需求。由于它的高效性,TF-IDF 模型在搜索引擎等实际应用中被广泛使用。

以下是本人使用Python实现该算法的思路，如有不当之处望各位大牛指导一二。

二、TF-IDF算法概述

关于TF-IDF算法的描述网上很多，我就不拾前人牙慧了，感兴趣的筒子们可参考这篇通俗易懂的文章：TF-IDF模型的概率解释。这篇文章中给出了很多数学公式，但数学的美妙在于其每个符号在现实中都是有着极其和谐的对应关系的。在下文中，我将用更通俗的方法阐述下个人的理解。

对于一篇文档来说，它与关键字 w[i] 的相关度取决于它包含的所有词中该关键词的频率。这其实挺直观的，一篇文档中包含关键词w[i]越多，那么它与关键字w[i]相关度也就越大。但是，如果仅仅取关键词的频数的话，那么比较长的文档包含该关键词的频数很可能远远大于比较短的文档的。因而为了协调文档长度的影响，相关度的衡量应取关键词w[i]占文档总词数的频率。

那有多个关键词的话该怎样衡量一篇文档出现的情况该怎样衡量文档的综合相关度呢？最简单的当然是把它们都加起来，但这样一来新的问题又出现了。假设某个关键词w[j]出现在很多篇文档里，另一关键词w[k]仅在一小部分文档(记为集合U[i])里出现，那按照常理来说是不是匹配了更多关键词的文档集U[k]与给出的搜索关键词w[k],w[j]相关度更大？鉴于这种情况，我们需要给每一篇文档包含的每一个关键词的相关度加一个权值。《TF-IDF模型的概率解释》这篇文章里给出了该权值的推导过程。依本人的肤浅理解，这个权值为关键词w在所有文档集中所蕴含的信息熵。

这样TF-IDF算法的模型就出来了：

TF-IDF (q, d) = sum { i = 1..k | TF (w[i], d) *IDF(w[i]) }

IDF为逆向文档频率：

IDF = log (n / docs (w, D))

TF表示词条在文档d中出现的频率：

TF (w,d)= count (w, d) / sum { i = 1..n| count (w, d[i]) }

三、算法实现

1、文档预处理

获取了足够多的文档后，需要对文档进行预处理，以加快搜索的速度。=由于linux系统和windows系统的默认编码不同，Python在处理中文文档时可能会出错，所以也需要对不同编码格式的文档预处理成同一编码格式的文档。因而在文档预处理这一模块需要有以下几个步骤：读取文档 -> 删除不需要的字符（如回车符\n、制表符\t、空格等）-> 转换成unicode格式 -> 对文档分词 -> 转换成utf-8格式写入txt文档。

这些步骤的实现主要使用了以下几个模块

字符串修剪模块str_replace.py

这个模块就一个函数，代码如下：

def str_replace(str_source,char,*words):  
    str_temp=str_source      
    for word in words:  
        str_temp=str_temp.replace(word,char)  
    return str_temp

str_replace(str_source,char,*words)函数接受两个或两个以上的参数，str_source是需要处理的字符串，char是要替换的目标字符，words是可变字符串的元组，对字符串str_source中的每一个出现在words里的字符均替换成统一字符char。

在主程序里可以这样使用str_replace(content_temp,"","\t","\n","")，即将content_temp里的每一个"\t","\n",""字符都删掉。

字符串格式转换模块StrToUni.py

这个模块主要有两个函数StrToUni_try和StrToUni。

由于输入字符串的格式可能没法实现知道，因而需要进行unicode解码的尝试，在解码尝试成功后再进行转码。这一过程分两个步骤，StrToUni_try和StrToUni两个函数分别完成。StrToUni_try函数主要负责判断字符串是不是某一格式，这个函数返回字符串的正确编码格式。StrToUni函数负责使用StrToUni_try返回的编码格式将字符串转化成unicode格式。代码如下：

def StrToUni_try(str,type_1):  
    try:  
        str.decode(type_1)  
    except UnicodeDecodeError:  
        return False  
    else:  
        return True  
def StrToUni(str,*type_list):  
    if not type_list:  
        if StrToUni_try(str,'utf-8'):  
            return str.decode('utf-8')  
        else:  
            print "输入的源文件的编码格式不是utf-8"  
    else:  
        for type_2 in type_list:  
            if StrToUni_try(str,type_2):  
                return str.decode(type_2)  
            else:  
                if type_2==type_list[-1]:  
                    print "输入的源文件的编码格式不在您提供的格式列表中"

StrToUni_try函数尝试解码的格式列表可以在GrobalParament.py模块里设置。用户需要单独使用时可以手动输入，如:

StrToUni(str,”utf-8”,”GBK”,”GB2312”)

分词模块full_word_cut.py和half_word_cut.py

在本程序中使用结巴分词对中文文档分词。在本程序中有两种搜索模式，一是全文搜索，二是关键词搜索。前一种的实现方法是使用结巴分词将中文文档进行全文分词再存储，后者的实现方法是使用结巴分词提取关键词再存储。由full_word_cut.py和half_word_cut.py模块分别实现这两种分词功能。

Unicode格式转字符串模块UniToStr.py

这一模块也是由两个函数组成，转换的过程与StrToUni.py模块类似。代码如下：

def UniToStr_try(str,type_1):  
    try:  
        str.encode(type_1)  
    except LookupError:  
        return False  
    else:  
        return True  
def UniToStr(str,*out_Format):  
    if not out_Format:  
        return str.encode('utf-8')  
    else:  
        for type_2 in out_Format:  
            if UniToStr_try(str,type_2):  
                return str.encode(type_2)  
            else:  
                if type_2==out_Format[-1]:  
                    print "输入的目标编码格式不正确"

该模块的目标编码格式可以在GrobalParament.py模块里设置，也可由用户手动设置。需要注意的是，这个模块只会根据目标编码格式列表里的第一个正确的编码格式进行编码，也就是说允许目标编码格式里的格式名是错误的格式名。如果用户不大确定某一编码格式叫什么名，可以多输入几个可能的组合作尝试。如用户不记得linux下的编码格式的准确名，可以在GrobalParament.py模块里修改OutputFormatList。

如OutputFormatList=[“utf-9”,”uft-8”,”utf-8”]，事实上国际标准里没有“utf-9”、”uft-8”两种编码格式，但”utf-8”是正确的，程序就能正确地输出”utf-8”格式的字符串。这样程序就有了一定的容错性。

文档预处理顶层模块prepro_file.py

这个模块需要使用上文所提到的五个模块，主要功能是将一个目录里的所有文档都处理完后按照给定的编码格式输出到指定的文档中。主要代码如下：

#这个函数用于预处理文件处理过程中采用unicode编码  
import os  
from str_replace import str_replace  
from TF_IDF.StrToUni import StrToUni  
import GrobalParament  
from full_word_cut import fullcut  
from half_word_cut import halfcut  
from UniToStr import UniToStr  
def prepro_file(fl_in_url,re_out_url,*wd_be_del):  
    in_url=fl_in_url.replace('\\','/')  
    out_url=re_out_url.replace('\\','/')  
    try:  
        try:  
            fl_in=os.listdir(in_url)  
        except WindowsError:  
            print "您输入的预处理文档目录有误"  
        try:  
            re_out=open(out_url,'w')  
        except WindowsError:  
            print "您输入的结果文档输出目录有误"  
    except NameError:  
        pass  
    else:  
        for file in fl_in:  
            afile_url=fl_in_url+'/'+file  
            if os.path.isfile(afile_url):  
                afile=open(afile_url,"r")  
                content_temp="".join(afile.readlines())  
                if not wd_be_del:  
                    content=str_replace(content_temp,"","\t","\n"," ") 
                                       #删除某些特殊字符如\t,\n等以保证是一行的连续的  
                else:  
                    content=str_replace(content_temp,'',*wd_be_del)  
                con_unicode=StrToUni(content,*(GrobalParament.InputFormatList))  
                if GrobalParament.pattern=="full":  
                    cut_result=fullcut(con_unicode)  
                else:  
                    cut_result=halfcut(con_unicode)  
                s_fl_Name=UniToStr(file,*(GrobalParament.OutputFormatList))  
                re_out.write(s_fl_Name+'\t')  
                key_word_out=[]  
                for key_word in cut_result:  
                    s_key_word=UniToStr(key_word,*(GrobalParament.OutputFormatList))  
                    key_word_out.append(s_key_word)  
                out_str=','.join(key_word_out)  
                re_out.write(out_str)  
                re_out.write('\n')

这一模块没啥好说的，主要是将前文所讲的模块统和到一起完成文档的预处理功能。需要注意的是，这里使用正则表达式滤除了结巴分词的结果集中只有一个字的词，如”的”，”之”，”地”等用处不大的词。

在本程序中，预处理的结果默认的输出格式为：

文件名\t分词结果\n

⋯ ⋯

最后的结果保存为:E:\\EclipseProjection\\Python\\TF_IDF\\test\\pro_res.txt

保存路径和保存文件名可以GrobalParament.py模块里设置，设置方法见下文。

2.TF-IDF算法模块

文档预处理后，就是本文的主菜TF-IDF算法模块了。以下我将详细的阐述TF-IDF算法的实现思路。

1) 打开前文得到的预处理文档pro_res.txt，并读取一行字符串（代表某一文档），赋给变量data；

2) 如果data不为空，则接步骤3)；

3) 将data拆分，用file_name记录文档名，列表data_temp_2为该文档的所有词语的列表。data_temp_len为该文档词语总数，同时代表文档总数的变量files_num加1；

4) 对于每一个关键词word，先判断此时file_name代表的文档的所有词语中是否有关键词word，如果有的话接5)；

5) 判断word是否在其他文档里出现，没有出现过得话在字典word_in_allfiles_stat新增一元素，以该关键字为键名，并赋键值为1；否则的话将字典word_in_allfiles_stat里键名为word对应的关键字的元素的键值加1；

6) 如果字典word_in_afile_stat没有以此时file_name变量指向的文档名的元素，即当前文档包含已经遍历过了的关键字，则新建一个元素，文件名为file_name变量指向的文档名，键值为一空字典；

7) 在6)中得到的空字典里新建一列表，保存file_name变量指向的文档中包含当前关键字的数量和当前文档的总词语数量；

8) 重复以上2）~ 7）知道读完整个预处理文档pro_res.txt；最后应得到两个字典word_in_allfiles_stat和word_in_afile_stat。这两个字典的结构分别如下：

word_in_allfiles_stat{A:n1,B:n2,⋯⋯}

word_in_afile_stat{a:{A:[aA,suma],B:[aB,suma],⋯⋯}, ⋯⋯}, b:{A:[bA,sumb],B:[bB,sumb],⋯⋯}, ⋯⋯},⋯⋯}

在以上两个字典中集合{A,B,⋯⋯}代表搜索关键字集，集合{a,b,⋯⋯}代表文档集，{aA,aB,⋯⋯}代表文档包含关键字的个数，{suma,sumb,⋯⋯}代表文档对应的总词数。

例如，在word_in_afile_stat中a:{A:[aA,suma],B:[aB,suma],⋯⋯}代表文档a中包含关键字A的词数为aA，包含关键字B的词数为aB⋯⋯，文件a的总词数为suma，若a没有关键词B，则字典中不包含B这个键名。

在word_in_allfiles_stat中，代表共有n1个文档包含关键词A，n2个文档包含关键词B。

9) 最后格局字典word_in_allfiles_stat和word_in_afile_stat里的内容计算最终结果，最终按照结果值的由高到低返回指定数目的文档名。

这部分的源代码如下：

from __future__ import division  
import math  
from TF_IDF import GrobalParament  
def TF_IDF_Compute(file_import_url_temp,file_export_url_temp,*words):  
    file_import_url=file_import_url_temp.replace('\\','/')  
    file_export_url=file_export_url_temp.replace('\\','/')  
    data_source=open(file_import_url,'r')  
    data=data_source.readline()  
    word_in_afile_stat={}  
    word_in_allfiles_stat={}  
    files_num=0  
    while(data!=""):  
        data_temp_1=[]  
        data_temp_2=[]  
        data_temp_1=data.split("\t") #file name and key words of a file       
        data_temp_2=data_temp_1[1].split(",")#key words of a file  
        file_name=data_temp_1[0]  
        data_temp_len=len(data_temp_2)  
        files_num+=1  
        for word in words:  
            if word in data_temp_2:  
                if not word_in_allfiles_stat.has_key(word):  
                    word_in_allfiles_stat[word]=1  
                else:  
                    word_in_allfiles_stat[word]+=1  
                  
                if not word_in_afile_stat.has_key(file_name):  
                    word_in_afile_stat[file_name]={}  
                if not word_in_afile_stat[file_name].has_key(word):  
                    word_in_afile_stat[file_name][word]=[]  
                    word_in_afile_stat[file_name][word].append(data_temp_2.count(word))  
                    word_in_afile_stat[file_name][word].append(data_temp_len)  
        data=data_source.readline()  
    data_source.close()  
    if (word_in_afile_stat) and (word_in_allfiles_stat) and (files_num !=0):  
        TF_IDF_result={}  
        for filename in word_in_afile_stat.keys():  
            TF_IDF_result[filename]={}  
            for word in word_in_afile_stat[filename].keys():  
                word_n=word_in_afile_stat[filename][word][0]  
                word_sum=word_in_afile_stat[filename][word][1]  
                with_word_sum=word_in_allfiles_stat[word]  
                TF_IDF_result[filename][word]=((word_n/word_sum))*(math.log10(files_num/with_word_sum))
        TF_IDF_total={}  
        for filename in TF_IDF_result.keys():  
            TF_IDF_total[filename]=reduce(lambda x,y:x+y,TF_IDF_result[filename].values())  
        result_temp=[]  
        result_temp=sorted(TF_IDF_total.iteritems(),key=lambda x:x[1],reverse=True)  
        k=GrobalParament.result_file_num  
        result=[]  
        for item in result_temp:  
            if k!=0:  
                result.append(item[0])  
                k-=1  
            else:  
                break         
    else:  
        result=["None"]  
    if GrobalParament.out_to_file:  
        export=open(file_export_url,'w')  
        for item in result:  
            export.write(item+'\n')  
        export.close()  
    else:  
        return result

3. 参数设置模块

该模块为GrobalParament.py，其参数对应的含义如表所示：

全局参数含义
参数	含义
InputFormatList	输入文件的编码列表
OutputFormatList	输出文件的编码列表
pattern	搜索模式:"full"为全文搜索模式,"keys"为关键词搜索模式
n	关键词搜索模式时的关键词数量
ruler_list	不需要的字符
result_file_num	需要查找多少个相关文档
out_to_file	是否需要输出为txt
PreprocessResultDir	预处理文件输出目录
PreprocessResultName	预处理输出文件名
ResultFileNameDir	搜索结果文件输出目录
ResultFileName	搜索结果文件名