[数据结构与算法]Leetcode--Java--677. 键值映射

题目描述

样例描述

示例 1：

输入：
["MapSum", "insert", "sum", "insert", "sum"]
[[], ["apple", 3], ["ap"], ["app", 2], ["ap"]]
输出：
[null, null, 3, null, 5]

解释：
MapSum mapSum = new MapSum();
mapSum.insert("apple", 3);  
mapSum.sum("ap");           // 返回 3 (apple = 3)
mapSum.insert("app", 2);    
mapSum.sum("ap");           // 返回 5 (apple + app = 3 + 2 = 5)

思路

字典树 + DFS

1. 将key挂在字典树上，然后在最后一个字母的下一个字母上存储val。因为遍历完prefix后，也会停留到最后一个字母的下一个。 而且如果是最后一个的 ，dfs求和时会多算上prefix最后一个字母的val
2. 对于前缀如果找不到，直接返回0不存在。否则就dfs遍历以当前前缀最后一个字母的所有孩子结点，只要不为空就往下遍历。然后累加val，不存在的话val为0
3. 求和部分要递归求和，因为p的孩子结点可能还有孩子

代码

class MapSum {
    
    class TrieNode {
        //先默认为0
        //存在key的话，在key的最后一个字母的后面一个上存储val
        int val = 0;
        TrieNode tns[];
        public TrieNode() {
            tns = new TrieNode[26];
        }
    }
    private TrieNode root;

    public MapSum() {
        root = new TrieNode();
    }
    
    public void insert(String key, int val) {
        TrieNode p = root;
        for (int i = 0; i < key.length(); i ++ ) {
            char c = key.charAt(i);
            int u = c - 'a';
            if (p.tns[u] == null) {
                p.tns[u] = new TrieNode();
            }
            //如果是key最后一个字母，存到孩子val上
            if (i == key.length() - 1) {
                p.tns[u].val = val;
            }
            p = p.tns[u];
        }
    }
    public int sum(String prefix) {
        TrieNode p = root;
        for (char c: prefix.toCharArray()) {
            int u = c - 'a';
            //前缀不存在，直接返回
            if (p.tns[u] == null) {
                return 0;
            }
            p = p.tns[u];
        }
        //前缀走完后，dfs暴力枚举当前字母的所有孩子，累加val（如果没有的话就是加0不影响）
        return dfs(p);
    }
    public int dfs(TrieNode root) {
        int ans = root.val;
        for (TrieNode p: root.tns) {
            if (p != null) {
                //递归，因为当前p也可能有孩子
                ans += dfs(p);
            }
        }
        return ans;
    }
}

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