Map Sum Pairs LeetCode solution in Python

1
class MapSum:
2

3
    def __init__(self):
4
        self.trie = {}
5

6
    def insert(self, key: str, val: int) -> None:
7
        node = self.trie
8
        for ch in key:
9
            node = node.setdefault(ch, {})
10
        node["val"] = val
11
        return
12

13
    def sum(self, prefix: str) -> int:
14
        def traverse(node):
15
            nonlocal res
16
            if not node:
17
                return
18
            if "val" in node:
19
                res += node["val"]
20
            for child in node:
21
                if child == "val":
22
                    continue
23
                traverse(node[child])
24
            return
25

26
        node = self.trie
27
        for ch in prefix:
28
            node = node.get(ch, {})
29
            if not node:
30
                return 0
31
        res = 0
32
        traverse(node)
33
        return res