Leetcode 677.鍵值映射（Map Sum Pairs）

Leetcode 677.鍵值映射

1 題目描述（Leetcode題目鏈接）

  實現一個 MapSum 類裏的兩個方法，insert 和 sum。

對於方法 insert，你將得到一對（字符串，整數）的鍵值對。字符串表示鍵，整數表示值。如果鍵已經存在，那麼原來的鍵值對將被替代成新的鍵值對。

對於方法 sum，你將得到一個表示前綴的字符串，你需要返回所有以該前綴開頭的鍵的值的總和。

輸入: insert("apple", 3), 輸出: Null
輸入: sum("ap"), 輸出: 3
輸入: insert("app", 2), 輸出: Null
輸入: sum("ap"), 輸出: 5

2 題解

  字典樹基本操作的實現，可以參照第208題，求和的過程就是一個深度優先搜索，首先找到前綴的最後一個節點，以這個節點爲根進行深度優先搜索就可以了。

class TreeNode:

    def __init__(self):
        self.children = collections.defaultdict(TreeNode)
        self.value = 0


class MapSum:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = TreeNode()
        

    def insert(self, key: str, val: int) -> None:
        node = self.root
        for char in key:
            node = node.children[char]
        node.value = val        

    def sum(self, prefix: str) -> int:
        node = self.root
        for char in prefix:
            node = node.children.get(char)
            if not node:
                return 0
        
        def dfs(node):
            res = node.value
            for char in node.children:
                res += dfs(node.children[char])
            return res
        
        return dfs(node)

        


# Your MapSum object will be instantiated and called as such:
# obj = MapSum()
# obj.insert(key,val)
# param_2 = obj.sum(prefix)