677. Map Sum Pairs

class MapSum {

    private static class Trie {

        private static class TrieNode {
            private final TrieNode[] children;
            private int val;

            public TrieNode() {
                this.children = new TrieNode[26];
            }
        }

        private final TrieNode root;

        public Trie() {
            this.root = new TrieNode();
        }

        public void insert(String key, int val) {
            TrieNode node = root;
            for (int i = 0; i < key.length(); i++) {
                int offset = key.charAt(i) - 'a';
                if (node.children[offset] == null) {
                    node.children[offset] = new TrieNode();
                }
                node = node.children[offset];
            }
            node.val = val;
        }

        public int sum(String prefix) {
            TrieNode node = root;
            for (int i = 0; i < prefix.length(); i++) {
                int offset = prefix.charAt(i) - 'a';
                if (node.children[offset] == null) {
                    return 0;
                }
                node = node.children[offset];
            }

            return dfs(node);
        }

        private int dfs(TrieNode node) {
            if (node == null) {
                return 0;
            }

            int val = node.val;
            for (TrieNode child : node.children) {
                val += dfs(child);
            }
            return val;
        }
    }

    private final Trie trie;

    public MapSum() {
        this.trie = new Trie();
    }

    public void insert(String key, int val) {
        trie.insert(key, val);
    }

    public int sum(String prefix) {
        return trie.sum(prefix);
    }

}

注意该题在相同 key 时会替换 val 的值，所以不能轻易在插入过程中累加 val 值，即使要累加，也需要引入 HashMap 等结构来处理相同 key 插入时的场景。

References

677. Map Sum Pairs
剑指 Offer II 066. 单词之和