master-algorithms-py/tries/map_sum.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# author: bt3gl

'''
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 str.
'''

class Node:
    
    def __init__(self):
        self.children = {}
        self.score = 0


class MapSum:

    def __init__(self):
        self.map = {}
        self.root = Node()

    def insert(self, key: str, val: int) -> None:
        delta = val - self.map.get(key, 0)
        self.map[key] = val
        
        node = self.root
        node.score += delta
        for c in key:
            node = node.children.setdefault(c, Node())
            node.score += delta

    def sum(self, prefix: str) -> int:
        node = self.root
        for c in prefix:
            if c not in node.children:
                return 0
            node = node.children[c]
          
        return node.score