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

题目描述

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

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

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

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

示例 1：

输入：n = 7, edges = [[0,1],[1,2],[3,1],[4,0],[0,5],[5,6]], restricted = [4,5]
输出：4
解释：上图所示正是这棵树。
在不访问受限节点的前提下，只有节点 [0,1,2,3] 可以从节点 0 到达。

示例 2：

输入：n = 7, edges = [[0,1],[0,2],[0,5],[0,4],[3,2],[6,5]], restricted = [4,2,1]
输出：3
解释：上图所示正是这棵树。
在不访问受限节点的前提下，只有节点 [0,5,6] 可以从节点 0 到达。

解题思路

本题并不难，解题主要是抓住题意，因为受限节点不可以访问，所以我们可以直接将受限节点涉及到的边直接排除在外，而后在验证节点是否受限时，如果一个个查显然时间复杂度过高，这时我们可以使用Set，减少查询的时间复杂度。而后进行一次dfs就可以了，而后我们还需要知道，因为这是一棵树，所以节点不会重复访问，所以我们直接++即可。

代码如下

class Solution {
    int cnt=0;
    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        Set<Integer> set=new HashSet<Integer>();
        List<Integer> lists[]=new ArrayList[n];
        for(int i:restricted)//存入set
            set.add(i);
        for(int i=0;i<n;i++)
            lists[i]=new ArrayList<>();
        for(int i=0;i<n-1;i++){
            int x=edges[i][0];
            int y=edges[i][1];
            if(set.contains(x)||set.contains(y))//不进行边加入
                continue;
            lists[x].add(y);
            lists[y].add(x);
        }
        boolean flag[]=new boolean[n];
        flag[0]=true;
        dfs(0,lists,flag);
        return cnt;
    }

    public void dfs(int p,List<Integer> lists[],boolean flag[]){
        cnt++;//不会重复直接++
        List<Integer> list=lists[p];
        for(int i=0;i<list.size();i++){
            int l=list.get(i);
            if(!flag[l]){
                flag[l]=true;
                dfs(l,lists,flag);
                flag[l]=false;
            }
        }
    }
}