Design a map that allows you to do the following:
Implement the MapSum
class:
MapSum()
Initializes the MapSum
object.void insert(String key, int val)
Inserts the key-val
pair into the map. If the key
already existed, the original key-value
pair will be overridden to the new one.int sum(string prefix)
Returns the sum of all the pairs' value whose key
starts with the prefix
.
Example 1:
Input ["MapSum", "insert", "sum", "insert", "sum"] [[], ["apple", 3], ["ap"], ["app", 2], ["ap"]] Output [null, null, 3, null, 5] Explanation MapSum mapSum = new MapSum(); mapSum.insert("apple", 3); mapSum.sum("ap"); // return 3 (apple = 3) mapSum.insert("app", 2); mapSum.sum("ap"); // return 5 (apple + app = 3 + 2 = 5)
Constraints:
1 <= key.length, prefix.length <= 50
key
and prefix
consist of only lowercase English letters.1 <= val <= 1000
50
calls will be made to insert
and sum
.This problem requires designing a data structure that efficiently handles two operations: inserting key-value pairs and calculating the sum of values associated with keys sharing a common prefix. A Trie (prefix tree) is an ideal choice for this because it's optimized for prefix-based searches.
Approach:
The solution uses a Trie to store the keys and their associated values. Each node in the Trie represents a character in a key. Each node also stores a val
which represents the sum of all values associated with keys passing through that node.
insert(key, val)
: When inserting a key-value pair, the algorithm traverses the Trie, adding nodes as needed. For each node visited, the val
is incremented by the difference between the new value and the previous value associated with that key (if it existed). This efficiently handles updates.
sum(prefix)
: To find the sum of values for keys starting with a given prefix, the algorithm traverses the Trie down to the node representing the last character of the prefix. The val
of this node is returned, which represents the sum of all values associated with keys that share this prefix.
Time Complexity Analysis:
insert(key, val)
: The time complexity is O(m), where 'm' is the length of the key. This is because we traverse at most 'm' nodes in the Trie to insert a key.
sum(prefix)
: The time complexity is O(n), where 'n' is the length of the prefix. This is because we traverse at most 'n' nodes to find the sum of values for keys with that prefix.
Space Complexity Analysis:
The space complexity is O(NM), where N is the number of keys inserted and M is the average length of the keys. In the worst-case scenario where all keys share the same prefix, space complexity can be O(NM).
Code Explanation (Python3 as an Example):
The Python code implements the Trie data structure within the MapSum
class. The insert
method efficiently updates the Trie with the new key-value pair, adjusting node values accordingly. The sum
method quickly calculates the prefix sum by traversing the Trie to the prefix's end node.
class Trie:
def __init__(self):
self.children = [None] * 26 # Array to store children nodes (a-z)
self.val = 0 # Sum of values for keys passing through this node
def insert(self, w: str, x: int):
node = self
for c in w:
idx = ord(c) - ord('a') # Get index (0-25) for character
if node.children[idx] is None:
node.children[idx] = Trie()
node = node.children[idx]
node.val += x
def search(self, w: str) -> int:
node = self
for c in w:
idx = ord(c) - ord('a')
if node.children[idx] is None:
return 0
node = node.children[idx]
return node.val
class MapSum:
def __init__(self):
self.d = defaultdict(int) # To store key-value pairs for efficient updates
self.tree = Trie() # Trie for prefix sum calculations
def insert(self, key: str, val: int) -> None:
x = val - self.d[key] #Difference for efficient update
self.d[key] = val
self.tree.insert(key, x)
def sum(self, prefix: str) -> int:
return self.tree.search(prefix)
The other language implementations follow a similar structure, adapting the Trie and MapSum
class implementations to their respective syntax and data structures. The core logic and time/space complexities remain consistent across all languages.