题目描述

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

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

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

样例

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

算法

(Trie树) $O(L)$

1. 采用 Trie 树的数据结构，结点需要存储总和 sum_val, 以该结点为结尾的单词值 word_val，和26个 nxt 指针。
2. 插入时，需要先寻找树中是否已经存在了这个单词，若存在，则需要返回该单词的 word_val 以备更新字典树上的结点值。然后再将新单词的价值在树中的结点路径上累计，累计过程中，需要去掉原单词（如果有）的价值。
3. 对前缀查询和时，只需要返回路径上最后一个结点的 sum_val 即可；如果不存在，则直接返回 0。

时间复杂度

• 路径长度等于单词长度，故单次操作的时间就是遍历符合单词的路径，时间复杂度为 $O(L)$。

C++ 代码

class MapSum {
public:
    struct TreeNode {
        int sum_val, word_val;
        TreeNode *nxt[26];
        TreeNode() {
            sum_val = 0;
            word_val = 0;
            memset(nxt, NULL, sizeof(nxt));
        }
    };
    TreeNode *root;
    /** Initialize your data structure here. */
    MapSum() {
        root = new TreeNode();
    }

    int find(string key) {
        int l = key.length();
        TreeNode *p = root;
        for (int i = 0; i < l; i++) {
            if (p -> nxt[key[i] - 'a'] == NULL)
                return 0;
            p = p -> nxt[key[i] - 'a'];
        }
        return p -> word_val;
    }

    void insert(string key, int val) {
        int original_val = find(key);
        int l = key.length();
        TreeNode *p = root;
        for (int i = 0; i < l; i++) {
            if (p -> nxt[key[i] - 'a'] == NULL)
                p -> nxt[key[i] - 'a'] = new TreeNode();
            p = p -> nxt[key[i] - 'a'];

            p -> sum_val -= original_val;
            p -> sum_val += val;
        }
        p -> word_val = val;
    }

    int sum(string prefix) {
        int l = prefix.length();
        TreeNode *p = root;
        for (int i = 0; i < l; i++) {
            if (p -> nxt[prefix[i] - 'a'] == NULL)
                return 0;
            p = p -> nxt[prefix[i] - 'a'];
        }
        return p -> sum_val;
    }
};

/**
 * Your MapSum object will be instantiated and called as such:
 * MapSum obj = new MapSum();
 * obj.insert(key,val);
 * int param_2 = obj.sum(prefix);
 */