06_TF-IDF算法代码示例

import numpy as np # 数值计算、矩阵运算、向量运算
import pandas as pd # 数值分析、科学计算

# 定义文档
docA = 'The cat sat on my bed'
docB = 'The dog sat on my knees'

# 切割文档
bowA = docA.split(' ')
bowB = docB.split(' ')
# bowA # ['The', 'cat', 'sat', 'on', 'my', 'bed']
# bowB # ['The', 'dog', 'sat', 'on', 'my', 'knees']

# 构建词库
wordSet = set(bowA).union(set(bowB))
# wordSet # {'The', 'bed', 'cat', 'dog', 'knees', 'my', 'on', 'sat'}

# 用字典来保存词出现的次数
wordDictA = dict.fromkeys(wordSet, 0)
wordDictB = dict.fromkeys(wordSet, 0)
wordDictA
wordDictB

# 遍历文档，统计词数
for word in bowA:
    wordDictA[word] += 1
for word in bowB:
    wordDictB[word] += 1

pd.DataFrame([wordDictA, wordDictB])

def computeTF(wordDict, bow):
    # 用一个字典对象保存 TF，把所有对应于 bow 文档里的 TF都计算出来
    tfDict = {}
    nbowCount = len(bow)

    for word, count in wordDict.items():
        tfDict[word] = count / nbowCount
    return tfDict

# 测试
tfA = computeTF(wordDictA, bowA)
tfB = computeTF(wordDictB, bowB)
print(tfA)
print(tfB)

    {'on': 0.16666666666666666, 'dog': 0.0, 'cat': 0.16666666666666666, 'The': 0.16666666666666666, 'knees': 0.0, 'my': 0.16666666666666666, 'sat': 0.16666666666666666, 'bed': 0.16666666666666666}
    {'on': 0.16666666666666666, 'dog': 0.16666666666666666, 'cat': 0.0, 'The': 0.16666666666666666, 'knees': 0.16666666666666666, 'my': 0.16666666666666666, 'sat': 0.16666666666666666, 'bed': 0.0}

def computeIDF(wordDictList):
    # 用一个字典对象保存 IDF，每个词作为 key，初始值为 0
    idfDict = dict.fromkeys(wordDictList[0], 0)
    # 总文档数量
    N = len(wordDictList)
    import math

    for wordDict in wordDictList:
        # 遍历字典中的每个词汇，统计 Ni
        for word, count in wordDict.items():
            if count > 0 :
                # 先把 Ni 增加 1，存入到 idfDict 中
                idfDict[word] += 1

    # 已经得到所有词汇 i 对应的 Ni，现在根据公式把它替换成 idf 值
    for word, Ni in idfDict.items():
        idfDict[word] = math.log10((N + 1)/(Ni + 1))
    return idfDict

# 测试
idfs = computeIDF([wordDictA, wordDictB])
idfs

    {'on': 0.0,
     'dog': 0.17609125905568124,
     'cat': 0.17609125905568124,
     'The': 0.0,
     'knees': 0.17609125905568124,
     'my': 0.0,
     'sat': 0.0,
     'bed': 0.17609125905568124}

def computeTFIDF(tf, idfs):
    tfidf = {}
    for word, tfval in tf.items():
        tfidf[word] = tfval * idfs[word]
    return tfidf

# 测试
tfidfA = computeTFIDF(tfA, idfs)
tfidfB = computeTFIDF(tfB, idfs)

pd.DataFrame([tfidfA, tfidfB])