朴素贝叶斯垃圾邮件文档识别代码

下面是一个简单的朴素贝叶斯垃圾邮件文档识别代码的示例：

import os
import math

# 训练集和测试集所在目录
train_dir = './train'
test_dir = './test'

# 垃圾邮件和非垃圾邮件的概率
spam_prob = 0
ham_prob = 0

# 垃圾邮件和非垃圾邮件中每个单词的概率
spam_word_prob = {}
ham_word_prob = {}

# 统计训练集中垃圾邮件和非垃圾邮件的数量
spam_num = 0
ham_num = 0

# 统计训练集中垃圾邮件和非垃圾邮件中每个单词的数量
spam_word_num = {}
ham_word_num = {}

# 读取训练集中的邮件，并统计数量
for filename in os.listdir(train_dir):
    with open(os.path.join(train_dir, filename), 'r') as f:
        content = f.read()
        if 'spmsg' in filename:
            spam_num += 1
            for word in content.split():
                if word not in spam_word_num:
                    spam_word_num[word] = 0
                spam_word_num[word] += 1
        else:
            ham_num += 1
            for word in content.split():
                if word not in ham_word_num:
                    ham_word_num[word] = 0
                ham_word_num[word] += 1

# 计算垃圾邮件和非垃圾邮件的概率
spam_prob = spam_num / (spam_num + ham_num)
ham_prob = ham_num / (spam_num + ham_num)

# 计算垃圾邮件和非垃圾邮件中每个单词的概率
for word in spam_word_num:
    spam_word_prob[word] = (spam_word_num[word] + 1) / (sum(spam_word_num.values()) + len(spam_word_num))
for word in ham_word_num:
    ham_word_prob[word] = (ham_word_num[word] + 1) / (sum(ham_word_num.values()) + len(ham_word_num))

# 读取测试集中的邮件，并进行分类
correct_num = 0
total_num = 0
for filename in os.listdir(test_dir):
    with open(os.path.join(test_dir, filename), 'r') as f:
        content = f.read()
        spam_score = math.log(spam_prob)
        ham_score = math.log(ham_prob)
        for word in content.split():
            if word in spam_word_prob:
                spam_score += math.log(spam_word_prob[word])
            else:
                spam_score += math.log(1 / (sum(spam_word_num.values()) + len(spam_word_num)))
            if word in ham_word_prob:
                ham_score += math.log(ham_word_prob[word])
            else:
                ham_score += math.log(1 / (sum(ham_word_num.values()) + len(ham_word_num)))
        if spam_score > ham_score:
            if 'spmsg' in filename:
                correct_num += 1
            print(filename, 'spam')
        else:
            if 'spmsg' not in filename:
                correct_num += 1
            print(filename, 'ham')
        total_num += 1

# 输出分类准确率
print('Accuracy:', correct_num / total_num)

在这个代码中，我们首先读取训练集中的邮件，并统计垃圾邮件和非垃圾邮件的数量，以及每个单词在垃圾邮件和非垃圾邮件中出现的数量。然后，我们根据这些统计信息计算垃圾邮件和非垃圾邮件的概率，以及每个单词在垃圾邮件和非垃圾邮件中出现的概率。最后，我们读取测试集中的邮件，并根据朴素贝叶斯算法进行分类。在分类时，我们计算每个邮件属于垃圾邮件和非垃圾邮件的分数，然后根据分数来判断邮件的类别。最后，我们输出分类准确率。