Leetcode 2581. Count Number of Possible Root Nodes

Question

Leetcode 2581. Count Number of Possible Root Nodes

Woodyiiiiiii opened this issue a year ago · comments

这是一道re-rooting类型的树的题目。类似有：

re-rooting树题目大概是一个不确定根节点的多叉树，从所有可能的根节点为基础出题。

解法说到底就是两遍DFS。第一遍DFS类似预先处理，一般以0为根节点，然后找到0为根节点时的特殊解，记录缓存。然后观察从0逐步向子节点扩散，以这些字节点为根节点时，跟第一次DFS预先得到的对比，找到规律，从而优化时间复杂度。

class Solution {
    int ans = 0;

    public int rootCount(int[][] edges, int[][] guesses, int k) {
        int n = edges.length + 1;
        // convert edges to map
        Map<Integer, Set<Integer>> map = new HashMap<>();
        for (int[] edge : edges) {
            int u = edge[0], v = edge[1];
            map.putIfAbsent(u, new HashSet<>());
            map.putIfAbsent(v, new HashSet<>());
            map.get(u).add(v);
            map.get(v).add(u);
        }
        // convert guesses to map
        Map<Integer, Set<Integer>> guessMap = new HashMap<>();
        for (int[] g : guesses) {
            int u = g[0], v = g[1];
            guessMap.putIfAbsent(u, new HashSet<>());
            guessMap.get(u).add(v);
        }
        // start from 0
        int[] parent = new int[n];
        Arrays.fill(parent, -1);
        Queue<Integer> queue = new LinkedList<>();
        queue.offer(0);
        parent[0] = 0;
        while (!queue.isEmpty()) {
            int u = queue.poll();
            for (int v : map.get(u)) {
                if (parent[v] != -1) continue;
                parent[v] = u;
                queue.offer(v);
            }
        }
        // check if guess is correct when root is 0
        int cnt = 0;
        for (int[] g : guesses) {
            int u = g[0], v = g[1];
            if (parent[v] == u) cnt++;
        }
        if (cnt >= k) ans++;

        // re-rooting/dfs
        for (int c : map.get(0)) {
            dfs(c, 0, map, cnt, k, guessMap);
        }

        return ans;
    }

    private void dfs(int cur, int p, Map<Integer, Set<Integer>> map, int cnt, int k, Map<Integer, Set<Integer>> guessMap) {
        if (guessMap.containsKey(cur) && guessMap.get(cur).contains(p)) cnt++;
        if (guessMap.containsKey(p) && guessMap.get(p).contains(cur)) cnt--;
        if (cnt >= k) ans++;
        if (!map.containsKey(cur)) return;
        for (int c : map.get(cur)) {
            if (c == p) continue;
            dfs(c, cur, map, cnt, k, guessMap);
        }
    }
}