diff --git a/tries/map_sum.py b/tries/map_sum.py new file mode 100644 index 0000000..36efcaa --- /dev/null +++ b/tries/map_sum.py @@ -0,0 +1,42 @@ +#!/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 +