677. Map Sum Pairs

重新开始刷题的第1天。

今天的题目是677. Map Sum Pairs。

一道前缀树相关的题目。

一开始没看到题目中的“如果键已经存在，那么原来的键值对将被替代成新的键值对”。想当然的在实现插入时，直接增加值了。然后在第二个测例就错了。

class TireNode {
public:
    TireNode(int _value):child(256, nullptr), value(_value) {}

    vector<TireNode *> child;
    int value;
};

class MapSum {
public:
    /** Initialize your data structure here. */
    MapSum() {
        root = new TireNode(0);
    }
    
    void insert(string key, int val) {
        TireNode *p = root;
        for(int i = 0;i < key.size(); i++) {
            int index = key[i];
            if (p->child[index] == nullptr) {
                p->child[index] = new TireNode(0);
            }
            p = p->child[index];
            p->value += val;
        }
    }
    
    int sum(string prefix) {
        TireNode *p = root;
        for(int i = 0;i < prefix.size() && p != nullptr; i++) {
            int index = prefix[i];
            p = p->child[index];
        }
        return p == nullptr ? 0 : p->value;
    }
    TireNode *root;
};

因为这里希望能够更新值，所以最简单的修改方式就是给TireNode增加一个isLeaf的属性来标识该节点是否是某个键的末尾，然后只有当isLeaf == true时，节点的value才是有意义的。这样修改的话，插入部分不需要做太多修改，但是前缀和的求解却需要递归遍历子树。

class TireNode {
public:
    TireNode(int _value = 0):child(256, nullptr), value(_value), isLeaf(false) {}

    vector<TireNode *> child;
    int value;
    bool isLeaf;
};

class MapSum {
public:
    /** Initialize your data structure here. */
    MapSum() {
        root = new TireNode(0);
    }
    
    void insert(string key, int val) {
        TireNode *p = root;
        for(int i = 0;i < key.size(); i++) {
            int index = key[i];
            if (p->child[index] == nullptr) {
                p->child[index] = new TireNode();
            }
            p = p->child[index];
        }
        p->value  = val;
        p->isLeaf = true;
    }
    
    int sum(string prefix) {
        TireNode *p = root;
        for(int i = 0;i < prefix.size() && p != nullptr; i++) {
            int index = prefix[i];
            p = p->child[index];
        }
        return p == nullptr ? 0 : sum(p);
    }
    int sum(TireNode *p) {
        if (p == nullptr) return 0;
        int res = p->isLeaf ? p->value : 0;
        for(int i = 0;i < 256;i++) {
            res += sum(p->child[i]);
        }
        return res;
    }
    TireNode *root;
};

这样过是过了，但是速度贼慢，主要在于sum这里每次都需要重新求解。因此为了优化速度，我们给TrieNode结构增加一个sum属性，这个属性会在insert中维护，使其为该节点的前缀和。在维护sum属性的方式是通过在insert时，修改完末尾节点的value后，计算出新的value与旧的value（默认为0）的差值，然后向上更新。同时，为了加快速度和减少空间，这里还将TireNode的child从vector修改成unordered_map。

class TireNode {
public:
    TireNode(int _value = 0):value(_value), isLeaf(false), sum(0) {}

    unordered_map<char, TireNode *> child;
    int value, sum;
    bool isLeaf;
};

class MapSum {
public:
    /** Initialize your data structure here. */
    MapSum() {
        root = new TireNode(0);
    }
    
    void insert(string key, int val) {
        TireNode *p = root;
        stack<TireNode*> st;
        for(int i = 0;i < key.size(); i++) {
            st.push(p);
            char index = key[i];
            if (p->child[index] == nullptr) {
                p->child[index] = new TireNode();
            }
            p = p->child[index];
        }
        int diff = val - p->value;
        while(!st.empty()) {
            st.top()->sum += diff;
            st.pop();
        }
        p->sum += diff;
        p->value  = val;
        p->isLeaf = true;
    }
    
    int sum(string prefix) {
        TireNode *p = root;
        for(int i = 0;i < prefix.size() && p != nullptr; i++) {
            char index = prefix[i];
            p = p->child[index];
        }
        return p == nullptr ? 0 : p->sum;
    }
    int sum(TireNode *p) {
        if (p == nullptr) return 0;
        int res = p->isLeaf ? p->value : 0;
        for(auto &pair: p->child) {
            res += sum(pair.second);
        }
        return res;
    }
    TireNode *root;
};