1. 程式人生 > 實用技巧 >b_nk_實現Trie樹(用陣列記錄孩子結點)

b_nk_實現Trie樹(用陣列記錄孩子結點)

實現一個 Trie (字首樹),包含 insert, delete, isExist, 和 prefixCount 這四個操作。

class Node:
    def __init__(self, val) -> None:
        self.child=[None for i in range(26)]
        self.isEnd=False
        self.val=val
        self.cnt=0 #以結點val為尾巴的路徑數
class Trie:
    def __init__(self) -> None:
        self.root=Node('/')
    """
    新增word
    """
    def insert(self, word: str) -> None:
        cur=self.root
        for c in word:
            u=ord(c)-ord('a')
            if not cur.child[u]:
                cur.child[u]=Node(c)
            cur=cur.child[u]
            cur.cnt+=1
        cur.isEnd=True
    """
    刪除word
    """
    def delete(self, word: str) -> None:
        cur=self.root
        for c in word:
            u=ord(c)-ord('a')
            if not cur.child[u] or cur.child[u].cnt==0: 
                cur.child[u]=None
                return
            cur.child[u].cnt-=1
            cur=cur.child[u]
        cur.isEnd=False
    """
    查詢word是否在字典樹中
    """
    def isExist(self, word: str) -> str:
        cur=self.root
        for c in word:
            u=ord(c)-ord('a')
            if not cur.child[u]:
                return "NO"
            cur=cur.child[u]
        return "YES" if cur.isEnd else "NO"
    """
    返回以word為字首的單詞數量
    """
    def prefixCount(self, word: str) -> int:
        cur=self.root
        for c in word:
            u=ord(c)-ord('a')
            if not cur.child[u]: return 0
            cur=cur.child[u]
        return cur.cnt
    
class Solution:
    def trieU(self , ops ):
        trie,ans=Trie(),[]
        for op in ops:
            if op[0]=="1":   #新增
                trie.insert(op[1])
            elif op[0]=="2": #刪除
                trie.delete(op[1])
            elif op[0]=="3": #查詢是否存在
                ans.append(trie.isExist(op[1]))
            else:            #返回以op[1]開頭的單詞的數量
                ans.append(str(trie.prefixCount(op[1])))
        return ans