JZX轻语：简
LeetCode 677 - 键值映射
发表于2024年04月09日
Trie的经典题目，需要在Trie节点维护一个用于存储该前缀和的变量prefixSum，然后搜索的时候遍历到指定前缀的最后一个节点时，返回该变量即可。
需要注意的是，该题目允许insert已存在的单词，且会覆盖原来的val，因此，上述情况还需要更新所有前缀的prefixSum。为解决上述问题，还需要在Trie节点上新增一个变量finishedVal以存储以该节点结尾的单词val(如果该变量为0，意味仅存在相应的前缀，但没有对应的单词)，当insert时发现最后一个节点的finishedVal不为0，则需要依次更新所遍历节点的前缀和prefixSum。
class Trie:
def __init__(self):
self.children = {}
self.prefixSum = 0
# 以该节点结尾的单词val(因为会存在重复insert的情况, 需要覆盖掉之前的val以及更新前缀和)
# 如果该值不为0, 则意味着到该节点是有一个完整的单词的
self.finishedVal = 0
class MapSum:
def __init__(self):
self._root = Trie()
def insert(self, key: str, val: int) -> None:
visited_nodes = [self._root]
cur_node = self._root
for i, ch in enumerate(key):
if ch not in cur_node.children:
cur_node.children[ch] = Trie()
cur_node = cur_node.children[ch]
visited_nodes.append(cur_node)
old_val = visited_nodes[-1].finishedVal
if old_val != 0: # 先前存在相同的单词
# 更新所遍历节点的prefixSum
for node in visited_nodes:
node.prefixSum += (val - old_val)
else:
for node in visited_nodes:
node.prefixSum += val
visited_nodes[-1].finishedVal = val
def sum(self, prefix: str) -> int:
cur_node = self._root
for ch in prefix:
if ch not in cur_node.children:
return 0 # 没有该前缀
cur_node = cur_node.children[ch]
return cur_node.prefixSum