TF-IDF算法原理详解与Python代码实现

# -*- coding: utf-8 -*-

import math
import jieba
from collections import Counter
from collections import defaultdict


def TFIDF(dataset):
    # 统计每个文章中每个词语的tf
    j = 0
    doc_len = len(dataset)
    cntr = [0] * doc_len          # Counter每个文章
    total = [0] * doc_len          # 获取每个文章中各自的词语总数
    word_tf = [0] * doc_len          # 获取每个文章中每个词语的tf
    for sent in dataset:
        iters = jieba.cut(sent)
        cntr[j] = Counter(iters)    
        total[j] = sum(cntr[j].values())    
        word_tf[j] = defaultdict(int)
        for i in cntr[j].keys():        
            word_tf[j][i] = cntr[j][i] / total[j]
        j += 1
    
    
    # 构建词汇表和总词频
    word_cnt = defaultdict(int)
    for word_list in word_tf:
        for word in word_list:
            word_cnt[word] += 1
    
    # 统计词汇在文档中的频率
    doc_tf = defaultdict(int)
    for word in word_cnt:
        for article in dataset:
            if word in article:
                doc_tf[word] += 1
    
    # 构建词汇的idf
    word_idf = {}
    for word in doc_tf:
        word_idf[word] = math.log(doc_len / (doc_tf[word] + 1))
        
    # 构建每篇文章中词的tf-idf     
    t = 0    
    word_tf_idf = [0] * doc_len        
    for article in word_tf:
        word_tf_idf[t] = defaultdict(int)
        for word in article:
            word_tf_idf[t][word] = article[word] * word_idf[word]            
        t += 1
    
    # 对每篇文章每个词的tf-idf进行由大到小排序
    t = 0
    for i in word_tf_idf:
        for j in i:  
            word_tf[t] = dict(sorted(word_tf[t].items(), key=lambda x:x[1], reverse=True))
            word_tf_idf[t] = dict(sorted(word_tf_idf[t].items(), key=lambda x:x[1], reverse=True))        
        t += 1
            
    return word_tf_idf, word_tf
    
    

if __name__ == '__main__':

    dataset = ["中国的北京是中国的首都","那里有一句话叫做“北京欢迎您”,北京啊北京",
               "北京欢迎您！相信梦想，梦，相信有梦就会有奇迹"]    
    word_tf_idf,word_tf = TFIDF(dataset)
    
    print("=============tf-idf=============")
    for i in word_tf_idf:
        print(i)
    print("=============tf=============")
    for i in word_tf:
        print(i)