Create map_sum.py

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
+