動機

沒看清楚功能要求吃了3個WA

Problem

Design a map that allows you to do the following:

  • Maps a string key to a given value.
  • Returns the sum of the values that have a key with a prefix equal to a given string.

Implement the MapSum class:

  • MapSum() Initializes the MapSum object.
  • void insert(String key, int val) Inserts the key-val pair into the map. If the key already existed, the original key-value pair will be overridden to the new one.
  • int sum(string prefix) Returns the sum of all the pairs' value whose key starts with the prefix.

 

Example 1:

Input[MapSum, insert, sum, insert, sum][[], [apple, 3], [ap], [app, 2], [ap]]Output[null, null, 3, null, 5]ExplanationMapSum mapSum = new MapSum();mapSum.insert(apple, 3);  mapSum.sum(ap);           // return 3 (apple = 3)mapSum.insert(app, 2);    mapSum.sum(ap);           // return 5 (apple + app = 3 + 2 = 5)

 

Constraints:

  • 1 <= key.length, prefix.length <= 50
  • key and prefix consist of only lowercase English letters.
  • 1 <= val <= 1000
  • At most 50 calls will be made to insert and sum.

Sol

就是Trie，但是要有特別的dfs

class Trie:
    def __init__(self):
        self.chars = {}
        self.val = 0
    def insert(self,s,v):
        if not s:
            self.val = v
        else:
            if s[0] not in self.chars:
                self.chars[s[0]] = Trie()
            self.chars[s[0]].insert(s[1:], v)
    
    def pxsum(self,px,acc=0):
        if not px:
            return acc+sum([self.chars[k].pxsum("")+self.chars[k].val for k in self.chars],0)
        elif px[0] in self.chars:
            return self.chars[px[0]].pxsum(px[1:],acc+self.chars[px[0]].val if len(px) == 1 else 0)
        else:
            return 0

class MapSum:

    def __init__(self):
        self.trie = Trie()

    def insert(self, key: str, val: int) -> None:
        self.trie.insert(key,val)

    def sum(self, prefix: str) -> int:
        return self.trie.pxsum(prefix)