677 Map Sum Pairs

Implement a MapSum class withinsert, andsummethods.

For the methodinsert, you'll be given a pair of (string, integer). The string represents the key and the integer represents the value. If the key already existed, then the original key-value pair will be overridden to the new one.

For the methodsum, you'll be given a string representing the prefix, and you need to return the sum of all the pairs' value whose key starts with the prefix.

Example 1:

Input: insert("apple", 3), Output: Null
Input: sum("ap"), Output: 3
Input: insert("app", 2), Output: Null
Input: sum("ap"), Output: 5

The Idea: The datastructure that I ended up using is a plain old, non-optimized Trie. The basic idea is that for inserts, we just iterate down the trie using the prefix string, either creating new nodes along the way, or selecting a new path. At the end of the path, we modify the value of the node, and mark it as the end of string. My trie supported a single operation insert, for this behavior. With sum, we first iterate down to the specified prefix. We can do this just by following the path of the prefix string through the trie. If the prefix does not exist, we can detect this and return 0. Finally, we perform a DFS on the subtree the iteration left off at. This DFS will accumulate all the values of the subtree and return it. This is correct behavior because every node with the DFS tree is guaranteed to carry the prefix of the root.

Complexity: O(|S|) time for insert and O(N) time for sum, where N is the total number of nodes, which at worst case, scales linearly with the total number of characters inserted (but peaks at a maximum of 26 characters/branches per node). Same with space.

struct Node {
    Node(int val) : end_of_word(false), value(val) {}
    unordered_map<char, Node*> children;
    bool end_of_word;
    int value;
};

class Trie {
public:
    Trie() { _root = new Node(0); }
    void add(const string& name, int val) {
        Node *iter = _root;
        for (char c : name) {
            if (iter->children.find(c) == iter->children.end()) {
                iter->children.insert(make_pair(c, new Node(0)));
            }
            iter = iter->children.find(c)->second;
        }
        iter->value = val;
        iter->end_of_word = true;
        //print_trie(); cout << endl;
    }

    void print_trie() {
        Node *iter = _root;
        queue<Node*> q;
        q.push(_root);

        while (!q.empty()) {
            iter = q.front();
            q.pop();

            for (auto i : iter->children)
                cout << i.first << ' ';
            cout << endl;

            for (auto i : iter->children) {
                q.push(i.second);
            }
        }
    }

    Node *_root;
};


class MapSum {
public:
    /** Initialize your data structure here. */
    MapSum() {
        trie = Trie();
    }

    void insert(string key, int val) {
        trie.add(key, val);
    }

    int sum(string prefix) {
        int total_sum = 0;
        Node *start = iter_to_prefix(prefix);
        if (!start) return 0;
        _sum(start, total_sum);
        return total_sum;
    }

private:
    Trie trie;

    Node* iter_to_prefix(const string &prefix) {
        Node *root = trie._root;
        for (char c : prefix) {
            if (root->children.find(c) != root->children.end())
                root = root->children[c];
            else return nullptr;
        }
        return root;
    }

    void _sum(Node *cur, int &total_sum) {
        if (!cur->children.empty()) {
            total_sum += cur->value;
            for (auto i : cur->children) {
                _sum(i.second, total_sum);
            }
        }
        else total_sum += cur->value;
    }

};

Testing

int main() {
    // t1
    MapSum mp;
    mp.insert("apple", 3);
    cout << mp.sum("ap") << endl; // 3
    mp.insert("app", 2);
    cout << mp.sum("ap") << endl; // 5

    // t2
    MapSum mp2;
    mp2.insert("a", 3);
    cout << mp2.sum("ap") << endl; // 0 (not found)
    mp2.insert("b", 2);
    cout << mp2.sum("a") << endl; // 3
}