Trie 前缀树 Python实现

题目描述

字典树又称为前缀树或者Trie树，是处理字符串常用的数据结构。假设组成所有单词的字符仅是‘a’～‘z’，请实现字典树的结构，并包含以下四个主要的功能。void insert(String word)：添加word，可重复添加；void delete(String word)：删除word，如果word添加过多次，仅删除一次；boolean search(String word)：查询word是否在字典树中出现过(完整的出现过，前缀式不算)；int prefixNumber(String pre)：返回以字符串pre作为前缀的单词数量。现在给定一个m，表示有m次操作，每次操作都为以上四种操作之一。每次操作会给定一个整数op和一个字符串word，op代表一个操作码，如果op为1，则代表添加word，op为2则代表删除word，op为3则代表查询word是否在字典树中，op为4代表返回以word为前缀的单词数量（数据保证不会删除不存在的word）。

对于每次操作，如果op为3时，如果word在字典树中，请输出“YES”，否则输出“NO”；如果op为4时，请输出返回以word为前缀的单词数量，其它情况不输出。

输入： [[“1”,“qwer”],[“1”,“qwe”],[“3”,“qwer”],[“4”,“q”],[“2”,“qwer”],[“3”,“qwer”],[“4”,“q”]]

返回值：[“YES”,“2”,“NO”,“1”]

代码

class Trie: 
    def __init__(self):
        self.root = {} # 使用dict代替树结构
        
    def insert(self, word):
        if len(word) > 0:
            cur = self.root
            for c in word:
                if c not in cur:
                    cur[c] = {}
                cur = cur[c]
            cur['#'] = True
    
    def delete(self, word):
        if len(word) > 0 and self.find(word):
            stack = []
            cur = self.root
            is_end = True
            for c in word:
                stack.append(cur) # 不要写成append(cur[c])，我们添加的是从root到倒数第二个节点
                cur = cur[c]
                
            if len(cur) > 1: #判断这个单词是不是最后的词，即不是其他单词的前缀
                is_end = False
                
            for i, c in enumerate(word[::-1]):
                parent = stack.pop()
                if is_end == False: # 注意parent和parent[c]的区别
                    parent[c].pop('#')
                    break
                elif '#' in parent[c] and i == 0: 
                    parent.pop(c)
                elif len(parent[c]) == 0: # 只有当这条路径中只有这个词的时候，才能删掉这个字母
                    parent.pop(c)
                
               
            
            
    def search(self, word):
        if self.find(word):
            return 'Yes'
        else:
            return 'No'
        
    def find(self, word):
        if len(word) > 0:
            cur = self.root
            for c in word:
                if c not in cur:
                    return False
                cur = cur[c]
            if '#' in cur:
                return True
        
        return False
    
    def prefixNumber(self, pre):
        cur = self.root
        for c in pre:
            if c not in cur:
                return 0
            cur = cur[c]
        
        return self.countPaths(cur)
    
    def countPaths(self, cur):     
        if len(cur) == 1 and '#' in cur:
            return 1
        
        count = 0
        for c in cur:
            if c == '#':
                count += 1
                continue
            count += self.countPaths(cur[c]) #不要写成self.countPaths(c), c只是一个key，我们要传的是值
            
        return count
        

trie = Trie()
trie.insert('ap')
trie.insert('app')
trie.insert('apd')
trie.insert('appb')
print(trie.prefixNumber('ap'))
trie.delete('appb')
print(trie.search('ap'))
print(trie.search('app'))
print(trie.search('apd'))
print(trie.search('appb'))
print(trie.prefixNumber('ap'))

运行结果

4
Yes
Yes
Yes
No
3