题目来源
题目描述
题目解析
题目难点
MapSum *a = new MapSum();
a->insert("apple", 3);
std::cout << a->sum("ap") << "\n";
a->insert("app", 2);
a->insert("apple", 2); // 也就是说之前的apple要更新为2
std::cout << a->sum("ap") << "\n"; // 4
思路
- 首先字典树处理出所有的前缀以及前缀终点对应的val值
- 然后查询到前缀终点入口,再用dfs遍历所有的节点,统计所有绑定有val的值之和
实现
class MapSum {
struct TrieNode{
int val;
bool isFinished;
std::vector<TrieNode *> children;
TrieNode() : val(0), children(26, NULL){
}
};
TrieNode *root;
int getTotal(TrieNode *node){
if(node == nullptr){
return 0;
}
int ans = 0;
if(node->isFinished){
ans += node->val;
}
for (int i = 0; i < 26; ++i) {
if(node->children[i] != nullptr){
ans += getTotal(node->children[i]);
}
}
return ans;
}
public:
MapSum() {
root = new TrieNode();
}
void insert(string key, int val) {
TrieNode *curr = root;
for(auto ch : key){
int idx = ch - 'a';
if(curr->children[idx] == nullptr){
curr->children[idx] = new TrieNode();
}
curr = curr->children[idx];
}
curr->val = val;
curr->isFinished = true;
}
int sum(string prefix) {
TrieNode *curr = root;
for(auto ch : prefix){
int idx = ch - 'a';
if(curr->children[idx] == nullptr){
return 0;
}
curr = curr->children[idx];
}
return getTotal(curr);
}
};