题解 | 字典树的实现

import sys


class TrieNode:
    def __init__(self):
        self.nexts = {}
        self.e = 0  # 记录以该节点结尾的单词个数
        self.p = 0  # 记录经过该节点的单词总数


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

    def insert(self, word: str) -> None:
        node = self.root
        node.p += 1
        for char in word:
            if char not in node.nexts:
                node.nexts[char] = TrieNode()
            node = node.nexts[char]
            node.p += 1
        node.e += 1

    # 查询前缀树里word单词出现了几次
    def countWordsEqualTo(self, word: str) -> int:
        node = self.root
        for char in word:
            if char not in node.nexts:
                return 0
            node = node.nexts[char]
        return node.e

    # 查询前缀树里以prefix为前缀的单词出现了几次
    def countWordsStartingWith(self, prefix: str) -> int:
        node = self.root
        for char in prefix:
            if char not in node.nexts:
                return 0
            node = node.nexts[char]
        return node.p

    # 如果此前word插入过，删掉一次
    def erase(self, word: str) -> None:
        if self.countWordsEqualTo(word) > 0:
            node = self.root
            node.p -= 1
            for char in word:
                node.nexts[char].p -= 1
                if node.nexts[char].p == 0:
                    del node.nexts[char]
                    return
                node = node.nexts[char]
            node.e -= 1

my_trie = Trie()

for line in sys.stdin:
    a = line.split()
    if len(a) > 1:
        op, word = a[0], a[1]
        if op == '1':
            my_trie.insert(word)
        if op == '2':
            my_trie.erase(word)
        if op == '3':
            if my_trie.countWordsEqualTo(word) > 0:
                print("YES")
            else:
                print("NO")
        if op == '4':
            print(my_trie.countWordsStartingWith(word))