用Python实现树形结构的Bloom Filter算法

import mmh3
import bitarray

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_word = False

class BloomFilter:
    def __init__(self, capacity=1000000, error_rate=0.001):
        self.capacity = capacity
        self.error_rate = error_rate
        self.bit_array = bitarray.bitarray(capacity)
        self.bit_array.setall(0)
        self.hash_count = self.get_hash_count()
        self.trie_root = TrieNode()

    def add(self, word):
        if self.lookup(word):
            return
        self._add_to_trie(word)
        for i in range(self.hash_count):
            hash_code = mmh3.hash(word, i) % self.capacity
            self.bit_array[hash_code] = 1

    def lookup(self, word):
        if not self._lookup_in_trie(word):
            return False
        for i in range(self.hash_count):
            hash_code = mmh3.hash(word, i) % self.capacity
            if not self.bit_array[hash_code]:
                return False
        return True

    def get_hash_count(self):
        m = self.capacity
        k = -(m / (self.error_rate * math.log(1.0 - math.exp(-math.log(2.0)))))
        return int(k)

    def _add_to_trie(self, word):
        curr = self.trie_root
        for c in word:
            if c not in curr.children:
                curr.children[c] = TrieNode()
            curr = curr.children[c]
        curr.is_word = True

    def _lookup_in_trie(self, word):
        curr = self.trie_root
        for c in word:
            if c not in curr.children:
                return False
            curr = curr.children[c]
        return curr.is_word

if __name__ == '__main__':
    bf = BloomFilter(capacity=1000000, error_rate=0.001)
    bf.add("pidancode.com")
    bf.add("皮蛋编程")
    print(bf.lookup("pidancode.com"))
    print(bf.lookup("pianodengcheng"))