2368.受限条件下可到达节点的数目

目标

现有一棵由 n 个节点组成的无向树，节点编号从 0 到 n - 1 ，共有 n - 1 条边。

给你一个二维整数数组 edges ，长度为 n - 1 ，其中 edges[i] = [ai, bi] 表示树中节点 ai 和 bi 之间存在一条边。另给你一个整数数组 restricted 表示 受限 节点。

在不访问受限节点的前提下，返回你可以从节点 0 到达的 最多 节点数目。

注意，节点 0 不 会标记为受限节点。

思路

自然的想法是构建图，将受限节点从中删除，然后深度优先遍历，同时记录节点个数。这里构建的图主要是为了获取其连通节点进行dfs，HashSet不太适合。因为数据可能并不是连续存储的，要先计算元素的Hash值，然后从桶中取出链表或者红黑树，才能找到元素。在本例中，性能会下降一倍。

代码

/**
 * @date 2024-03-02 15:39
 */
public class ReachableNodes {
    public int res = 1;
    boolean[] isRestricted;

    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        List<Integer>[] g = new ArrayList[edges.length + 1];
        isRestricted = new boolean[edges.length + 1];
        for (int i : restricted) {
            isRestricted[i] = true;
        }
        for (int i = 0; i < g.length; i++) {
            g[i] = new ArrayList<>();
        }
        for (int[] edge : edges) {
            if (isRestricted[edge[0]] || isRestricted[edge[1]]) {
                continue;
            }
            g[edge[0]].add(edge[1]);
            g[edge[1]].add(edge[0]);
        }
        dfs(0, -1, g);
        return res;
    }

    public void dfs(int root, int parent, List<Integer>[] g) {
        for (Integer n : g[root]) {
            if (n == parent) {
                continue;
            }
            res++;
            dfs(n, root, g);
        }
    }
}

性能

看了官网的答案还可以使用并查集，耗时只要10ms，有时间可以看看。