Simple Trie, Python


  • -1
    Trie = lambda: collections.defaultdict(Trie) 
    class MapSum(object):
    
        def __init__(self):
            self.trie = Trie()
            self.seen = {}
            
        def insert(self, key, val):
            cur = self.trie
            for c in key:
                cur = cur[c]
                if cur[None] and key not in self.seen:
                    cur[None] += val
                elif key in self.seen:
                    cur[None] = cur[None] - self.seen[key] + val
                else:
                    cur[None] = val
            self.seen[key] = val
            
        def sum(self, prefix):
            cur = self.trie
            for c in prefix:
                cur = cur[c]
            if cur[None]:
                return cur[None]
            return 0
    

  • 1
    J

    Failed, should update prefix sum
    ["MapSum", "insert", "sum", "insert", "sum","insert","sum"]
    [[], ["apple",3], ["ap"], ["app",2], ["ap"],["apple",4],["ap"]]

    @TianQ said in Simple Trie, Python:

    Trie = lambda: collections.defaultdict(Trie)
    class MapSum(object):

    def __init__(self):
        self.trie = Trie()
        self.seen = set()
        
    def insert(self, key, val):
        cur = self.trie
        for c in key:
            cur = cur[c]
            if cur[None] and key not in self.seen:
                cur[None] += val
            else:
                cur[None] = val
        self.seen.add(key)
        
    def sum(self, prefix):
        cur = self.trie
        for c in prefix:
            cur = cur[c]
        if cur[None]:
            return cur[None]
        return 0

  • 0

    @jordandong
    Thanks for your comments. Yes you are right. Answer updated. Thanks again.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.