TF-IDF算法代码示例0.引入依赖1.定义数据和预处理2.进行词数统计3.计算词频 TF4.计算逆文档频率 IDF5.计算 TF-IDF
TF-IDF算法代码示例
0.引入依赖
import numpy as np # 数值计算、矩阵运算、向量运算
import pandas as pd # 数值分析、科学计算
1.定义数据和预处理
# 定义文档
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'}
2.进行词数统计
# 用字典来保存词出现的次数
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])
输出结果如下：
3.计算词频 TF
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}
4.计算逆文档频率 IDF
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}
5.计算 TF-IDF
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])
输出结果如下：
