Code & Func

重新开始刷题的第1天。

一道前缀树相关的题目。

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

1
class TireNode {
2
public:
3
    TireNode(int _value):child(256, nullptr), value(_value) {}
4

5
    vector<TireNode *> child;
6
    int value;
7
};
8

9
class MapSum {
10
public:
11
    /** Initialize your data structure here. */
12
    MapSum() {
13
        root = new TireNode(0);
14
    }
15

22 collapsed lines
16
    void insert(string key, int val) {
17
        TireNode *p = root;
18
        for(int i = 0;i < key.size(); i++) {
19
            int index = key[i];
20
            if (p->child[index] == nullptr) {
21
                p->child[index] = new TireNode(0);
22
            }
23
            p = p->child[index];
24
            p->value += val;
25
        }
26
    }
27

28
    int sum(string prefix) {
29
        TireNode *p = root;
30
        for(int i = 0;i < prefix.size() && p != nullptr; i++) {
31
            int index = prefix[i];
32
            p = p->child[index];
33
        }
34
        return p == nullptr ? 0 : p->value;
35
    }
36
    TireNode *root;
37
};

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

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

5
    vector<TireNode *> child;
6
    int value;
7
    bool isLeaf;
8
};
9

10
class MapSum {
11
public:
12
    /** Initialize your data structure here. */
13
    MapSum() {
14
        root = new TireNode(0);
15
    }
32 collapsed lines
16

17
    void insert(string key, int val) {
18
        TireNode *p = root;
19
        for(int i = 0;i < key.size(); i++) {
20
            int index = key[i];
21
            if (p->child[index] == nullptr) {
22
                p->child[index] = new TireNode();
23
            }
24
            p = p->child[index];
25
        }
26
        p->value  = val;
27
        p->isLeaf = true;
28
    }
29

30
    int sum(string prefix) {
31
        TireNode *p = root;
32
        for(int i = 0;i < prefix.size() && p != nullptr; i++) {
33
            int index = prefix[i];
34
            p = p->child[index];
35
        }
36
        return p == nullptr ? 0 : sum(p);
37
    }
38
    int sum(TireNode *p) {
39
        if (p == nullptr) return 0;
40
        int res = p->isLeaf ? p->value : 0;
41
        for(int i = 0;i < 256;i++) {
42
            res += sum(p->child[i]);
43
        }
44
        return res;
45
    }
46
    TireNode *root;
47
};

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

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

5
    unordered_map<char, TireNode *> child;
6
    int value, sum;
7
    bool isLeaf;
8
};
9

10
class MapSum {
11
public:
12
    /** Initialize your data structure here. */
13
    MapSum() {
14
        root = new TireNode(0);
15
    }
40 collapsed lines
16

17
    void insert(string key, int val) {
18
        TireNode *p = root;
19
        stack<TireNode*> st;
20
        for(int i = 0;i < key.size(); i++) {
21
            st.push(p);
22
            char index = key[i];
23
            if (p->child[index] == nullptr) {
24
                p->child[index] = new TireNode();
25
            }
26
            p = p->child[index];
27
        }
28
        int diff = val - p->value;
29
        while(!st.empty()) {
30
            st.top()->sum += diff;
31
            st.pop();
32
        }
33
        p->sum += diff;
34
        p->value  = val;
35
        p->isLeaf = true;
36
    }
37

38
    int sum(string prefix) {
39
        TireNode *p = root;
40
        for(int i = 0;i < prefix.size() && p != nullptr; i++) {
41
            char index = prefix[i];
42
            p = p->child[index];
43
        }
44
        return p == nullptr ? 0 : p->sum;
45
    }
46
    int sum(TireNode *p) {
47
        if (p == nullptr) return 0;
48
        int res = p->isLeaf ? p->value : 0;
49
        for(auto &pair: p->child) {
50
            res += sum(pair.second);
51
        }
52
        return res;
53
    }
54
    TireNode *root;
55
};