面试算法66：单词之和

题目

请设计实现一个类型MapSum，它有如下两个操作。

• 函数insert，输入一个字符串和一个整数，在数据集合中添加一个字符串及其对应的值。如果数据集合中已经包含该字符串，则将该字符串对应的值替换成新值。
• 函数sum，输入一个字符串，返回数据集合中所有以该字符串为前缀的字符串对应的值之和。

例如，第1次调用函数insert添加字符串"happy"和它的值3，此时如果输入"hap"调用函数sum则返回3。第2次调用函数insert添加字符串"happen"和它的值2，此时如果输入"hap"调用函数sum则返回5。

分析

首先定义前缀树节点的数据结构。由于每个字符串对应一个数值，因此需要在节点中增加一个整数字段。如果一个节点对应一个字符串的最后一个字符，那么该节点的整数字段的值就设为字符串的值；如果一个节点对应字符串的其他字符，那么该节点的整数字段将被设为0。由于这个题目只关注字符串对应的值之和，这些值已经在节点中的整数字段得以体现，因此节点中没有必要包含一个布尔变量标识节点是否对应字符串的最后一个字符。

解

public class MapSum {

    private TrieNode root;

    public MapSum() {
        root = new TrieNode();
    }

    public static void main(String[] args) {
        MapSum mapSum = new MapSum();
        mapSum.insert("happy", 3);
        System.out.println(mapSum.sum("hap"));
        mapSum.insert("happen", 2);
        System.out.println(mapSum.sum("hap"));
    }

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

            node = node.children[ch - 'a'];
        }

        node.value = val;
    }

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

            node = node.children[ch - 'a'];
        }

        return getSum(node);
    }

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

        int result = node.value;
        for (TrieNode child : node.children) {
            result += getSum(child);
        }

        return result;
    }

    class TrieNode {
        public TrieNode[] children;
        public int value;

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