Leetcode 2421. Number of Good Paths

[Woodyiiiiiii]

这道题是类似多叉树中求path的问题。

我总结了下，对于树/图问题，我的解决方法有如下几种：

1. 直接DFS/BFS(对于二维数组移动问题比较常见)
2. 树形DP问题(比如Binary Tree Camera / Maximum path sum of the tree)，可解决多叉树，但关键是要给定起始点
3. 并查集算法(针对图)

所以我一开始的想法，看到是Path问题，就想利用树状DP来解决：

1. 找到根节点
2. 然后利用后序遍历，返回Map解决问题，注意相邻问题
3. 循环中，利用相乘计算ans

这里有一些问题：

1. 首先这类问题根本没有所谓的根节点，因为每个点都可以叫做根节点。除非题目指定root或某个vertex。比如1443. Minimum Time to Collect All Apples in a Tree
2. 在递归函数中计算ans时，在边多的情况下不可避免地会超时

class Solution {
    
    int ans = 0;

    public int numberOfGoodPaths(int[] vals, int[][] edges) {
        // key: the index of the node, value: the value of the node
        Map<Integer, Integer> nodeValueMap = new HashMap<>();
        int n = vals.length;
        ans = n;
        for (int i = 0; i < vals.length; i++) {
            nodeValueMap.put(i, vals[i]);
        }
        // convert edges to graph
        Map<Integer, List<Integer>> graph = new HashMap<>();
        for (int[] edge : edges) {
            int from = edge[0];
            int to = edge[1];
            graph.computeIfAbsent(from, k -> new ArrayList<>()).add(to);
            graph.computeIfAbsent(to, k -> new ArrayList<>()).add(from);
        }

        // find root of the multi-tree
        // int root = 0;
        // for (int i = 0; i < n; i++) {
        //     int nodeCnt = getNodeCnt(i, graph, new HashSet<>());
        //     if (nodeCnt == n) {
        //         root = i;
        //         break;
        //     }
        // }

        // dfs
        dfs(graph, -1, 0, nodeValueMap);

        return ans;
    }

    private Map<Integer, Integer> dfs(Map<Integer, List<Integer>> graph, int parent, int cur, Map<Integer, Integer> nodeValueMap) {
        Map<Integer, Integer> curNodeValueMap = new HashMap<>();
        int val = nodeValueMap.get(cur);
        Map<Integer, Integer> cntMap = new HashMap<>();
        List<Map<Integer, Integer>> childrenNodeValueMapList = new ArrayList<>();
        curNodeValueMap.put(val, 1);
        if (graph.containsKey(cur)) {
            for (int next : graph.get(cur)) {
                if (next == parent) {
                    continue;
                }
                Map<Integer, Integer> nextNodeValueMap = dfs(graph, cur, next, nodeValueMap);
                childrenNodeValueMapList.add(nextNodeValueMap);
                for (int key : nextNodeValueMap.keySet()) {
                    if (key >= val) {
                        curNodeValueMap.put(key, curNodeValueMap.getOrDefault(key, 0) + nextNodeValueMap.get(key));
                        cntMap.put(key, cntMap.getOrDefault(key, 0) + nextNodeValueMap.get(key));
                    }
                    if (key == val) {
                        ans += nextNodeValueMap.get(key);
                    }
                }
            }
        }

        // update ans
        for (Map<Integer, Integer> childNodeValueMap : childrenNodeValueMapList) {
            for (int key : childNodeValueMap.keySet()) {
                if (key >= val) {
                    if (cntMap.containsKey(key) && cntMap.get(key) >= childNodeValueMap.get(key)) {
                        cntMap.put(key, cntMap.get(key) - childNodeValueMap.get(key));
                    }
                    ans += (childNodeValueMap.get(key) * cntMap.getOrDefault(key, 0));
                }
            }
        }
        return curNodeValueMap;
    }

    private int getNodeCnt(int node, Map<Integer, List<Integer>> graph, Set<Integer> visited) {
        visited.add(node);
        int cnt = 1;
        if (!graph.containsKey(node)) {
            return cnt;
        }
        for (int next : graph.get(node)) {
            if (!visited.contains(next)) {
                cnt += getNodeCnt(next, graph, visited);
            }
        }
        return cnt;
    }

}

感觉代码思路没什么问题，但卡在最后十几个test用例上，最后TLE了。

那么正解是什么呢？如何只用O(n)或者O(nlgn)的时间复杂度呢？

可以这么想，题目要求找到一个起始点value相同的节点，然后在路径中还要要求所有节点的value小于等于起始点value。是不是能将树/图分割成一个个群？即使用并查集算法，从每一个value相同的初始孤立的节点开始；然后问题是如果在并查集中如何merge节点呢？先从value小的开始，按顺序到大的，这样我们就只需要merge相邻节点即可；最后记录所有节点的群，及群的节点个数，求组成的path个数，遍历出现value的节点，用Map存储，key表示群的根节点，val表示同时为根节点相等的value相等的个数；最后我们知道每个群有多少个相同的value节点，然后用数学方法求子集，即count*(count+1)/2。

class Solution {
    
    public int numberOfGoodPaths(int[] vals, int[][] edges) {
        // convert edges to graph
        Map<Integer, List<Integer>> graph = new HashMap<>();
        for (int[] edge : edges) {
            int from = edge[0];
            int to = edge[1];
            graph.computeIfAbsent(from, k -> new ArrayList<>()).add(to);
            graph.computeIfAbsent(to, k -> new ArrayList<>()).add(from);
        }
        // valueMap, key: value, value: index
        //  ??????? why treemap?? —— connect
        TreeMap<Integer, List<Integer>> valueMap = new TreeMap<>();
        for (int i = 0; i < vals.length; i++) {
            valueMap.computeIfAbsent(vals[i], k -> new ArrayList<>()).add(i);
        }
        // union find
        UnionFind uf = new UnionFind(vals.length);

        int ans = 0;
        for (List<Integer> list : valueMap.values()) {
            Map<Integer, Integer> graphMap = new HashMap<>();
            for (int idx : list) {
                if (!graph.containsKey(idx)) {
                    continue;
                }
                for (int next : graph.get(idx)) {
                    if (vals[next] <= vals[idx]) {
                        uf.union_set(idx, next);
                    }
                }
            }
            for (int idx : list) {
                int root = uf.find(idx);
                graphMap.put(root, graphMap.getOrDefault(root, 0) + 1);
            }
            for (int count : graphMap.values()) {
                ans += count * (count + 1) / 2;
            }
        }

        return ans;
    }

}

class UnionFind {
    int[] parent;
    int[] rank;

    public UnionFind(int size) {
        parent = new int[size];
        for (int i = 0; i < size; i++)
            parent[i] = i;
        rank = new int[size];
    }

    public int find(int x) {
        if (parent[x] != x)
            parent[x] = find(parent[x]);
        return parent[x];
    }

    public void union_set(int x, int y) {
        int xRoot = find(x);
        int yRoot = find(y);
        if (xRoot == yRoot)
            return;
        if (rank[xRoot] < rank[yRoot]) {
            parent[xRoot] = yRoot;
            rank[yRoot] += rank[xRoot];
        } else {
            parent[yRoot] = xRoot;
            rank[xRoot] += rank[yRoot];
        }
    }
}

Tips:

1. 一定要注意遍历下一批节点时，一定要先用containsKey判断下
2. 这题还是很难想到用并查集算法的，但很有趣
3. 最后的int count : graphMap.values()，key不重要了，我们只需要计算在一个群中相同value节点的个数

同时，注意多叉树的path问题，可以用排序+并查集解决。

类似题目：

• 1697. Checking Existence of Edge Length Limited Paths

---

• 2421. Number of Good Paths
• Official Solution