实现一个 MapSum 类里的两个方法，insert 和 sum。

对于方法 insert，你将得到一对（字符串，整数）的键值对。字符串表示键，整数表示值。如果键已经存在，那么原来的键值对将被替代成新的键值对。

对于方法 sum，你将得到一个表示前缀的字符串，你需要返回所有以该前缀开头的键的值的总和。

示例 1:

输入: insert(“apple”, 3), 输出: Null
输入: sum(“ap”), 输出: 3
输入: insert(“app”, 2), 输出: Null
输入: sum(“ap”), 输出: 5

题目地址

中文版

英文版

 

代码实现

class TrieTree(object):
	def __init__(self):
		self.root = {}
	
	def insert(self,word,value):
		node = self.root
		for c in word:
			node = node.setdefault(c,{})
		node["#"] = value
		
	def prefix(self,prefix):
		node = self.root
		for c in prefix:
			if(c not in node):
				return None
			node = node
		return node
	
	def cacl_prefix_sum(self,node):
		if(not node):
			return 0
		self.total = 0
		def recursion(node):
			if(not node):
				return
			for key in node.keys():
				if(key=="#"):
					self.total += node[key]
				else:
					recursion(node[key])
		recursion(node)
		return self.total
			
	
class MapSum(object):
	def __init__(self):
		"""
		Initialize your data structure here.
		"""
		self.trie = TrieTree()
		

	def insert(self, key, val):
		"""
		:type key: str
		:type val: int
		:rtype: None
		"""
		self.trie.insert(key,val)
		

	def sum(self, prefix):
		"""
		:type prefix: str
		:rtype: int
		"""
		node = self.trie.prefix(prefix)
		return self.trie.cacl_prefix_sum(node)


# Your MapSum object will be instantiated and called as such:
# obj = MapSum()
# obj.insert(key,val)
# param_2 = obj.sum(prefix)
                                                                            

 

 

5 次阅读