insert(key: str, value: int): Set a given key’s value in the map. If the key already exists, overwrite the value. sum(prefix: str): Return the sum of all values of keys that begin with a given prefix. For example, you should be able to run the following code: mapsum.insert(“columnar”, 3) assert mapsum.sum(“col”) == 3 mapsum.insert(“column”, 2) assert mapsum.sum(“col”) == 5
Java
import java.util.HashMap;
import java.util.Map;
class TrieNode {
Map<Character, TrieNode> children;
int value;
public TrieNode() {
children = new HashMap<>();
value = 0;
}
}
public class PrefixMapSum {
TrieNode root;
Map<String, Integer> map;
public PrefixMapSum() {
root = new TrieNode();
map = new HashMap<>();
}
public void insert(String key, int value) {
int diff = value - map.getOrDefault(key, 0);
map.put(key, value);
TrieNode node = root;
node.value += diff;
for (char c : key.toCharArray()) {
node.children.putIfAbsent(c, new TrieNode());
node = node.children.get(c);
node.value += diff;
}
}
public int sum(String prefix) {
TrieNode node = root;
for (char c : prefix.toCharArray()) {
if (!node.children.containsKey(c)) {
return 0;
}
node = node.children.get(c);
}
return node.value;
}
}
Javascript
class TrieNode {
constructor() {
this.children = new Map();
this.value = 0;
}
}
class PrefixMapSum {
constructor() {
this.root = new TrieNode();
this.map = new Map();
}
insert(key, value) {
const diff = value - (this.map.get(key) || 0);
this.map.set(key, value);
let node = this.root;
node.value += diff;
for (const char of key) {
if (!node.children.has(char)) {
node.children.set(char, new TrieNode());
}
node = node.children.get(char);
node.value += diff;
}
}
sum(prefix) {
let node = this.root;
for (const char of prefix) {
if (!node.children.has(char)) {
return 0;
}
node = node.children.get(char);
}
return node.value;
}
}
Python
class TrieNode:
def __init__(self):
self.children = {}
self.value = 0
class PrefixMapSum:
def __init__(self):
self.root = TrieNode()
self.map = {}
def insert(self, key, value):
diff = value - self.map.get(key, 0)
self.map[key] = value
node = self.root
node.value += diff
for char in key:
if char not in node.children:
node.children[char] = TrieNode()
node = node.children[char]
node.value += diff
def sum(self, prefix):
node = self.root
for char in prefix:
if char not in node.children:
return 0
node = node.children[char]
return node.value
In this implementation, we use a Trie data structure to store the keys and values. The TrieNode
class represents a node in the Trie, containing a dictionary of children nodes and a value field that holds the sum of all values from its children nodes.
The PrefixMapSum
class utilizes the Trie to implement the required methods. The insert
method updates the value for a given key, while also updating the values of all relevant Trie nodes. The sum
method traverses the Trie based on the given prefix and returns the sum of values for all keys that start with that prefix.
You can test the functionality with the provided example code:
pythonCopy codemapsum = PrefixMapSum()
mapsum.insert("columnar", 3)
print(mapsum.sum("col")) # Output: 3
mapsum.insert("column", 2)
print(mapsum.sum("col")) # Output: 5