Analysis

Brute-force solution

Apparently, a straightforward brute-force solution would be to store every {key, value} pair into a HashSet. And once the sum(String prefix) function is called, we scan the whole HashSet’s key set, and use String.startswith() to select keys that satisfiy our prefix requirement. Then we add up all the values and return.

Complexity

For each insert operation, the time complexity is O(1). For each sum operation, the time complexity is O(m*n) where m is the total number of keys being stored and n is the length of input prefix.

The space used by the HashMap is liner to the size of all input keys and values.

Optimization1: Use a Trie to store the {key:value} relationship

Considering that we are compare prefix, a Trie (prefix tree) is naturally appropriate for this task. So for each insert operation, we can use a Trie to store the input keys and at the end Trie node of each key, we store the value in this node.

For every sum operation, we just start from the root of the Trie and looking for the first TrieNode that has a prefix matches the given prefix. And from this node, we perform a BFS to get every key in its child branch and sum up the values.

Complexity

Each insert operation will take O(n) time where n is the length of the input key.

Each sum operation will take O(m*k) time where m is the length of the prefix and k is the number of words start with this this prefix.

The space used by the Trie is liner to the size of the total input.

Implementation

A Java implementation is showed as following:

class MapSum {
    class TrieNode {
        char curr;
        TrieNode[] child;
        boolean isEnd;
        int value;
        
        TrieNode(char curr) {
            this.curr = curr;
            this.child = new TrieNode[26];
        }
    }
    
    TrieNode root;

    /** Initialize your data structure here. */
    public MapSum() {
        // setup a dummy root node
        this.root = new TrieNode(' ');
    }
    
    public void insert(String key, int val) {
        char[] charKey = key.toCharArray();
        TrieNode pointer = this.root;
        for (char c : charKey) {
            if (pointer.child[c - 'a'] == null) {
                TrieNode curr = new TrieNode(c);
                pointer.child[c - 'a'] = curr;
            }
            pointer = pointer.child[c - 'a'];
        }
        pointer.isEnd = true;
        pointer.value = val;
    }
    
    public int sum(String prefix) {
        int ret = 0;
        if (prefix == null || prefix.isEmpty()) {
            return ret;
        }
        char[] charPrefix = prefix.toCharArray();
        TrieNode pointer = this.root;
        for (char c : charPrefix) {
            if (pointer.child[c - 'a'] == null) {
                return ret;
            }
            
            pointer = pointer.child[c - 'a'];
        }
        
        Queue<TrieNode> queue = new LinkedList<TrieNode>();
        queue.offer(pointer);
        while (!queue.isEmpty()) {
            TrieNode curr = queue.poll();
            ret += curr.value;
            for (TrieNode t : curr.child) {
                if (t == null) {
                    continue;
                }
                queue.offer(t);
            }
        }
        
        return ret;
    }
}

Optimization2: Store the sum directly

Actually there is still some unnecessary work we have done in the above approach. Since we only cares about the prefix sum rather than the exact value that is related to this key, we do not need to store the values corresponding to each input key. Instead, we can store the sum of current prefix in the Trie. In this way, the insert is still the same with the above process, but during the insert process, rather than only store the value at the end TrieNode of current key, we store the sum of every prefix along the way.

Complexity

The time complexity of insert operation is O(n) where n is the length of the input key.

The time complexity of sum operation is O(m) where m is the length of the input prefix.

The space used is still linear to the total input size.

Implementation

A Java implementation is showed as following:

class MapSum {
    class TrieNode {
        TrieNode[] child = new TrieNode[26];
        int val;
    }
    
    HashMap<String, Integer> map;
    TrieNode root;

    /** Initialize your data structure here. */
    public MapSum() {
        this.map = new HashMap<>();
        this.root = new TrieNode();
    }
    
    public void insert(String key, int val) {
        int delta = val - map.getOrDefault(key, 0);
        map.put(key, val);
        TrieNode curr = this.root;
        curr.val += delta;
        for (char c : key.toCharArray()) {
            if (curr.child[c - 'a'] == null) {
                curr.child[c - 'a'] = new TrieNode();
                curr.child[c - 'a'].val = delta;
            } else {
                curr.child[c - 'a'].val += delta;
            }
            
            curr = curr.child[c - 'a'];
        }
    }
    
    public int sum(String prefix) {
        if (prefix == null || prefix.isEmpty()) {
            return 0;
        }
        TrieNode curr = this.root;
        for (char c : prefix.toCharArray()) {
            if (curr.child[c - 'a'] == null) {
                return 0;
            }
            curr = curr.child[c - 'a'];
        }
        
        return curr.val;
    }
}