forked from SamirPaulb/DSAlgo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
06. Map Sum Pairs.py
40 lines (34 loc) Β· 1.86 KB
/
06. Map Sum Pairs.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
# https://leetcode.com/problems/map-sum-pairs/
'''
When you see word prefix in problem formulation, you should definitely think about trie. What we need here is classical trie with one additional field: freq - frequency of each node. For example when we add word apple with val = 3 we need to add this 3 to all a, p, p, l, e nodes. Also we need to deal with this phrase: If the key already existed, the original key-value pair will be overridden to the new one, so we need to keep dictionary of pairs key - value to underastand if they already exist in our tree. Then we evaluate delta = val - self.dic.get(key, 0) is amount we need to change our values. For sum function we find node in our tree and if it is not there, return 0, if we can reach it, return frequency of this node.
'''
class TrieNode:
def __init__(self):
self.children = {}
self.prefixCount = 0
class MapSum:
def __init__(self):
self.root = TrieNode()
self.dic = {}
def insert(self, key: str, val: int) -> None:
delta = val
if key in self.dic: # key already existed, the original key-value pair will be overridden to the new one. And val - self.dic[key] does this thing
delta = val - self.dic[key]
self.dic[key] = val
cur = self.root
for c in key:
if c not in cur.children:
cur.children[c] = TrieNode()
cur = cur.children[c]
cur.prefixCount += delta
def sum(self, prefix: str) -> int:
cur = self.root
for c in prefix:
if c not in cur.children:
cur.children[c] = TrieNode()
cur = cur.children[c]
return cur.prefixCount
'''
Time complexity: O(m) to insert key of length m as well it is O(m) to evaluate sum for prefix of length m.
Space complexity: O(T) where T it total length of all words.
'''