字典树（前缀树）的实现

题目

字典树又称为前缀树或Trie树，是处理字符串常见的数据结构。假设组成所有单词的字符仅是’a’~'z’的小写字母，请实现字典树结构，包括如下功能：

• insert，添加word，可重复添加
• delete，删除word，如果word添加过多次，仅删除一个
• search，查询word是否在字典树中
• prefix_number，返回以字符串pre为前缀的单词数量

实现

import random
import string


class TrieNode():
    def __init__(self):
        self._end = 0
        self._path = 0
        self._map = [None for _ in range(26)]


class Trie():
    def __init__(self):
        self._root = TrieNode()

    def insert(self, s):
        if s is None:
            return

        node = self._root
        for c in s:
            index = ord(c) - ord('a')
            assert(0 <= index < 26)
            if node._map[index] is None:
                node._map[index] = TrieNode()
            node = node._map[index]
            node._path += 1

        node._end += 1

    def delete(self, s):
        if s is None:
            return

        if self.search(s) is False:
            return

        node = self._root
        for c in s:
            index = ord(c) - ord('a')
            assert(0 <= index < 26)
            assert(node._map[index] is not None)
            node = node._map[index]
            node._path -= 1

        node._end -= 1

    def search(self, s):
        if s is None:
            return False

        node = self._root
        for c in s:
            index = ord(c) - ord('a')
            assert(0 <= index < 26)
            if node._map[index] is None:
                return False
            node = node._map[index]

        return node._end != 0

    def prefix_number(self, prefix):
        if prefix is None:
            return 0

        node = self._root
        for c in prefix:
            index = ord(c) - ord('a')
            assert(0 <= index < 26)
            if node._map[index] is None:
                return 0
            node = node._map[index]

        return node._path

测试

def test_trie():
    trie = Trie()

    words = {}
    for _ in range(10000):
        length = random.randint(1, 5)
        s = ''.join([random.choice(string.ascii_lowercase)
                     for __ in range(length)])
        words[s] = 1 if s not in words else words[s]+1
        trie.insert(s)

    for _ in range(1000):
        length = random.randint(1, 3)
        s = ''.join([random.choice(string.ascii_lowercase)
                     for __ in range(length)])
        count = 0
        for w in words:
            if w.startswith(s):
                count += words.get(w, 0)
        assert(count == trie.prefix_number(s))

    for s in list(words):
        assert(trie.search(s) is True)
        trie.delete(s)
        words[s] -= 1

    for s in words:
        if words[s] == 0:
            assert(trie.search(s) is False)
        else:
            assert(trie.search(s) is True)

    print('done')


if __name__ == '__main__':
    test_trie()

测试2

牛客网-字典树的实现