[LeetCode] 677. 键值映射

[Animenzzzz]

题目描述：

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

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

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

示例 1:

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

解题思路：使用字典树，在每个节点上多加一个权重(初始值为0)，用来记录当前的值。插入key的时候，对所有节点的权重值都进行累加。加完之后，要再进行一次处理：如果当前key是之前已经插入过的key（判断条件：是否还有子节点，有：则是子串，无：则是已经插入过的前缀），则要把之前的值减回来，因为之前多加了。

C++解题：

class TrieNode{
public:
    TrieNode *child[26];
    bool isWordEnd;
    int weights;
    TrieNode():isWordEnd(false){
        weights = 0;
        for (int i = 0; i < 26; i++) child[i] = nullptr;
    }
};
class MapSum {
public:
    /** Initialize your data structure here. */
    MapSum() {
        
    }
    
    void insert(string key, int val) {
        TrieNode *p = root;
        int sum = 0;
        for(char cc:key){
            if(!p->child[cc-'a']){
                p->child[cc-'a'] = new TrieNode();
            }
            sum = p->child[cc-'a']->weights;
            p->child[cc-'a']->weights = val+p->child[cc-'a']->weights;
            p = p->child[cc-'a'];
        }
        //如果是之前插入过的前缀，要减回来
        int flag = 0;
        for (int i = 0; i < 26; i++)
        {
            if(p->child[i]){
                flag = 1;//说明当前的key是已经插入过的子串，而不是之前插入过的前缀
                break;
            }
        }
        if(!flag){
            p = root;
            for(char cc:key){
                p->child[cc-'a']->weights = p->child[cc-'a']->weights-sum;
                p = p->child[cc-'a'];
            }
        }
        
    }
    int sum(string prefix) {
        int result = 0;
        TrieNode *p = root;
        for(char cc:prefix){
            p = p->child[cc-'a'];
            if(!p) return 0;
            result = p->weights;
        }
        return result;
    }
private:
    TrieNode *root = new TrieNode();
};