如何使用 Python 堆实现朴素贝叶斯算法？

import heapq
from collections import defaultdict

class NaiveBayes:
    def __init__(self):
        self.classes = set()
        self.count = defaultdict(lambda: defaultdict(int))
        self.total_count = defaultdict(int)
        self.frequencies = defaultdict(float)

    def fit(self, data):
        for d in data:
            c = d[0]
            self.classes.add(c)
            for feature in d[1:]:
                self.count[c][feature] += 1
                self.total_count[c] += 1
                self.frequencies[feature] += 1
        # 计算每个特征在所有样本中出现的概率
        n = len(data)
        for feature in self.frequencies:
            self.frequencies[feature] /= n

    def predict(self, data):
        results = []
        for d in data:
            max_prob = -1
            max_class = None
            for c in self.classes:
                prob = 1.0
                for feature in d[1:]:
                    # 如果某个特征不在训练集中出现过，忽略这个特征
                    if feature not in self.count[c]:
                        continue
                    # 计算属于这个类别的概率
                    prob *= (self.count[c][feature] / self.total_count[c])
                # 加权计算概率
                prob *= (self.total_count[c] / len(data))
                # 计算先验概率
                prob /= (len(self.classes) * self.frequencies[feature])
                # 更新分类结果
                if prob > max_prob:
                    max_prob = prob
                    max_class = c
            results.append(max_class)
        return results

data = [
    ['spam', 'pidancode', 'com'],
    ['spam', '皮蛋', '编程'],
    ['ham', '算法', '堆'],
    ['ham', '排序', '堆'],
    ['ham', 'python', '堆']
]

nb = NaiveBayes()
nb.fit(data)

test_data = [
    ['pidancode', 'com'],
    ['排序', '堆'],
    ['计算机', '编程']
]
results = nb.predict(test_data)
print(results)

['spam', 'ham', None]