C++ Trie: Storing the Sum at Each Node


  • 0
    I
    struct Trie {
        unordered_map<char, Trie*> m;
        int sum = 0;
        bool end = false;
        int endValue = 0; // the value of the key
    };
    
    class MapSum {
        Trie* root = nullptr;
    public:
        /** Initialize your data structure here. */
        MapSum() {
            root = new Trie();
        }
        
        void insert(string key, int val) {
            int length = key.length();
            Trie* crt = root;
            for (int i = 0; i < length; i++) {
                if (crt->m[key[i]]) {
                    crt = crt->m[key[i]];
                } else {
                    crt->m[key[i]] = new Trie();
                    crt = crt->m[key[i]];
                }
            }
            
            // if the key doesn't exist
            if (!crt->end) {
                crt = root;
                for (int i = 0; i < length; i++) {
                    crt = crt->m[key[i]];
                    crt->sum += val;
                }
                crt->endValue = val;
                crt->end = true;
            } else { // if the key exists
                int preValue = crt->endValue;
                crt = root;
                for (int i = 0; i < length; i++) {
                    crt = crt->m[key[i]];
                    crt->sum -= preValue - val;
                }
            }
        }
        
        int sum(string prefix) {
            int result = 0;
            int length = prefix.length();
            Trie* crt = root;
            for (int i = 0; i < length; i++) {
                if (crt->m[prefix[i]]) {
                    crt = crt->m[prefix[i]];
                } else {
                    return 0;
                }
            }
            return crt->sum;
        }
    };
    
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.