-
Notifications
You must be signed in to change notification settings - Fork 892
/
problem_232.py
56 lines (42 loc) · 1.27 KB
/
problem_232.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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class Trie:
def __init__(self):
self.size = 0
self.value = 0
self.letter_map = dict()
def __repr__(self):
return "{}: {}".format(self.letter_map, self.value)
def add_word(self, word, value):
if not word:
self.value += value
return
letter = word[0]
sub_trie = None
if letter in self.letter_map:
sub_trie = self.letter_map[letter]
else:
sub_trie = Trie()
self.letter_map[letter] = sub_trie
self.size += 1
sub_trie.add_word(word[1:], value)
self.value = sum([x.value for x in self.letter_map.values()])
def find_prefix(self, word):
if not word:
return self.value
letter = word[0]
if letter in self.letter_map:
sub_trie = self.letter_map[letter]
return sub_trie.find_prefix(word[1:])
return 0
class PrefixMapSum:
def __init__(self, trie):
self.trie = trie
def insert(self, key, value):
self.trie.add_word(key, value)
def sum(self, prefix):
return self.trie.find_prefix(prefix)
# Tests
pms = PrefixMapSum(Trie())
pms.insert("columnar", 3)
assert pms.sum("col") == 3
pms.insert("column", 2)
assert pms.sum("col") == 5