题目
实现一个 MapSum
类里的两个方法,insert
和 sum
。
对于方法 insert
,你将得到一对(字符串,整数)的键值对。字符串表示键,整数表示值。如果键已经存在,那么原来的键值对将被替代成新的键值对。
对于方法 sum
,你将得到一个表示前缀的字符串,你需要返回所有以该前缀开头的键的值的总和。
示例 1:
输入: insert("apple", 3), 输出: Null
输入: sum("ap"), 输出: 3
输入: insert("app", 2), 输出: Null
输入: sum("ap"), 输出: 5
方案一(空间换时间)
class MapSum:
def __init__(self):
"""
Initialize your data structure here.
insert ("apple", 3) to self.lookup
{
"a": {
"sum": 3,
"p": {
"sum": 3,
"p": {
"sum": 3,
"l": {
"sum": 3,
"e": {
"sum": 3,
"is_word": True
}
}
}
}
}
}
"""
self.lookup = {}
self.keys = {} # "apple" -> 3
def insert(self, key: str, val: int) -> None:
node = self.lookup
if key not in self.keys:
self.keys[key] = val
for c in key:
if c not in node:
node[c] = {
"sum": val
}
else:
node[c]["sum"] += val
node = node[c]
else:
before_val = self.keys[key]
diff = val - before_val
for c in key:
node[c]["sum"] += diff
node = node[c]
def sum(self, prefix: str) -> int:
node = self.lookup
for c in prefix:
if c not in node:
return 0
node = node[c]
return node["sum"]
方案二
class MapSum:
def __init__(self):
"""
Initialize your data structure here.
insert ("apple", 3) to self.lookup
{
"a": {
"p": {
"p": {
"l": {
"e": {
"sum": 3,
}
}
}
}
}
}
then insert ("app", 2) to self.lookup
{
"a": {
"p": {
"p": {
"sum": 2,
"l": {
"e": {
"sum": 3,
}
}
}
}
}
}
"""
self.lookup = {}
def insert(self, key: str, val: int) -> None:
'''时间复杂度 O(n)'''
node = self.lookup
for c in key:
if c not in node:
node[c] = {}
node = node[c]
node["sum"] = val
def sum(self, prefix: str) -> int:
node = self.lookup
for c in prefix:
if c not in node:
return 0
node = node[c]
return self.count(node)
def count(self, root) -> int:
'''计算以 node 为根节点,下的 sum 的和'''
_sum = 0
for key in root:
if key == "sum":
_sum += root["sum"]
else:
_sum += self.count(root[key])
return _sum
原文
https://leetcode-cn.com/explore/learn/card/trie/167/practical-application-i/647/