Leetcode 677 Solution

class TrieNode {
public:
    char ch;
    int count;
    vector<TrieNode*> children;

    TrieNode(char _ch)
    {
        ch = _ch;
        count = 0;
        children.resize(256);
    }
};

class MapSum {
    TrieNode* root;

public:
    /** Initialize your data structure here. */
    MapSum() {
        root = new TrieNode(0);
    }

    void insert(string key, int val) {
        TrieNode* curr = root;

        for (auto ch: key)
        {
            if (!curr->children[ch])
                curr->children[ch] = new TrieNode(ch);

            curr = curr->children[ch];
        }

        curr->count = val;
    }

    int sum(string prefix) {
        TrieNode* curr = root;

        for (auto ch: prefix)
        {
            if (!curr->children[ch])
                return 0;
            curr = curr->children[ch];
        }

        return find_local_sum(curr);
    }

    int find_local_sum(TrieNode* node)
    {
        int s = 0;

        for (int i = 0; i < 256; i++)
            if (node->children[i])
                s += find_local_sum(node->children[i]);

        return s + node->count;
    }
};

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