Question

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

English Version

题目描述

现有一棵由 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 到达。

 

提示：

• 2 <= n <= 105
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• edges 表示一棵有效的树
• 1 <= restricted.length < n
• 1 <= restricted[i] < n
• restricted 中的所有值 互不相同

解法

方法一：DFS

我们首先根据给定的边构建一个邻接表 $g$，其中 $g[i]$ 表示与节点 $i$ 相邻的节点列表。然后我们定义一个哈希表 $vis$，用于记录受限节点或者已经访问过的节点，初始时将受限节点加入到 $vis$ 中。

接下来我们定义一个深度优先搜索函数 $dfs(i)$，表示从节点 $i$ 出发，可以到达的节点数。在 $dfs(i)$ 函数中，我们首先将节点 $i$ 加入到 $vis$ 中，然后遍历节点 $i$ 的相邻节点 $j$，如果 $j$ 不在 $vis$ 中，我们递归调用 $dfs(j)$，并将返回值加到结果中。

最后我们返回 $dfs(0)$ 即可。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 为节点数。

class Solution:
  def reachableNodes(
    self, n: int, edges: List[List[int]], restricted: List[int]
  ) -> int:
    def dfs(i: int) -> int:
      vis.add(i)
      return 1 + sum(j not in vis and dfs(j) for j in g[i])

    g = defaultdict(list)
    for a, b in edges:
      g[a].append(b)
      g[b].append(a)
    vis = set(restricted)
    return dfs(0)

class Solution {
  private List<Integer>[] g;
  private boolean[] vis;

  public int reachableNodes(int n, int[][] edges, int[] restricted) {
    g = new List[n];
    vis = new boolean[n];
    Arrays.setAll(g, k -> new ArrayList<>());
    for (var e : edges) {
      int a = e[0], b = e[1];
      g[a].add(b);
      g[b].add(a);
    }
    for (int i : restricted) {
      vis[i] = true;
    }
    return dfs(0);
  }

  private int dfs(int i) {
    vis[i] = true;
    int ans = 1;
    for (int j : g[i]) {
      if (!vis[j]) {
        ans += dfs(j);
      }
    }
    return ans;
  }
}

class Solution {
public:
  int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
    vector<int> g[n];
    vector<int> vis(n);
    for (auto& e : edges) {
      int a = e[0], b = e[1];
      g[a].emplace_back(b);
      g[b].emplace_back(a);
    }
    for (int i : restricted) {
      vis[i] = true;
    }
    function<int(int)> dfs = [&](int i) {
      vis[i] = true;
      int ans = 1;
      for (int j : g[i]) {
        if (!vis[j]) {
          ans += dfs(j);
        }
      }
      return ans;
    };
    return dfs(0);
  }
};

func reachableNodes(n int, edges [][]int, restricted []int) int {
  g := make([][]int, n)
  for _, e := range edges {
    a, b := e[0], e[1]
    g[a] = append(g[a], b)
    g[b] = append(g[b], a)
  }
  vis := make([]bool, n)
  for _, v := range restricted {
    vis[v] = true
  }
  ans := 0
  var dfs func(u int)
  dfs = func(u int) {
    if vis[u] {
      return
    }
    vis[u] = true
    ans++
    for _, v := range g[u] {
      dfs(v)
    }
  }
  dfs(0)
  return ans
}

function reachableNodes(n: number, edges: number[][], restricted: number[]): number {
  const vis: boolean[] = Array(n).fill(false);
  const g: number[][] = Array.from({ length: n }, () => []);
  for (const [a, b] of edges) {
    g[a].push(b);
    g[b].push(a);
  }
  for (const i of restricted) {
    vis[i] = true;
  }
  const dfs = (i: number): number => {
    vis[i] = true;
    let ans = 1;
    for (const j of g[i]) {
      if (!vis[j]) {
        ans += dfs(j);
      }
    }
    return ans;
  };
  return dfs(0);
}

方法二：BFS

与方法一类似，我们首先根据给定的边构建一个邻接表 $g$，然后定义一个哈希表 $vis$，用于记录受限节点或者已经访问过的节点，初始时将受限节点加入到 $vis$ 中。

接下来我们使用广度优先搜索遍历整个图，统计可以到达的节点数。我们定义一个队列 $q$，初始时将节点 $0$ 加入到 $q$ 中，并且将节点 $0$ 加入到 $vis$ 中。然后我们不断从 $q$ 中取出节点 $i$，累加答案，并将节点 $i$ 的相邻节点中未访问过的节点加入到 $q$ 中，并将这些节点加入到 $vis$ 中。

遍历结束后，返回答案即可。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 为节点数。

class Solution:
  def reachableNodes(
    self, n: int, edges: List[List[int]], restricted: List[int]
  ) -> int:
    g = defaultdict(list)
    for a, b in edges:
      g[a].append(b)
      g[b].append(a)
    vis = set(restricted + [0])
    q = deque([0])
    ans = 0
    while q:
      i = q.popleft()
      ans += 1
      for j in g[i]:
        if j not in vis:
          q.append(j)
          vis.add(j)
    return ans

class Solution {
  public int reachableNodes(int n, int[][] edges, int[] restricted) {
    List<Integer>[] g = new List[n];
    boolean[] vis = new boolean[n];
    Arrays.setAll(g, k -> new ArrayList<>());
    for (var e : edges) {
      int a = e[0], b = e[1];
      g[a].add(b);
      g[b].add(a);
    }
    for (int v : restricted) {
      vis[v] = true;
    }
    Deque<Integer> q = new ArrayDeque<>();
    q.offer(0);
    int ans = 0;
    for (vis[0] = true; !q.isEmpty(); ++ans) {
      int i = q.pollFirst();
      for (int j : g[i]) {
        if (!vis[j]) {
          q.offer(j);
          vis[j] = true;
        }
      }
    }
    return ans;
  }
}

class Solution {
public:
  int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
    vector<int> g[n];
    vector<int> vis(n);
    for (auto& e : edges) {
      int a = e[0], b = e[1];
      g[a].emplace_back(b);
      g[b].emplace_back(a);
    }
    for (int i : restricted) {
      vis[i] = true;
    }
    queue<int> q{{0}};
    int ans = 0;
    for (vis[0] = true; !q.empty(); ++ans) {
      int i = q.front();
      q.pop();
      for (int j : g[i]) {
        if (!vis[j]) {
          vis[j] = true;
          q.push(j);
        }
      }
    }
    return ans;
  }
};

func reachableNodes(n int, edges [][]int, restricted []int) (ans int) {
  g := make([][]int, n)
  vis := make([]bool, n)
  for _, e := range edges {
    a, b := e[0], e[1]
    g[a] = append(g[a], b)
    g[b] = append(g[b], a)
  }
  for _, i := range restricted {
    vis[i] = true
  }
  q := []int{0}
  for vis[0] = true; len(q) > 0; ans++ {
    i := q[0]
    q = q[1:]
    for _, j := range g[i] {
      if !vis[j] {
        vis[j] = true
        q = append(q, j)
      }
    }
  }
  return
}

function reachableNodes(n: number, edges: number[][], restricted: number[]): number {
  const vis: boolean[] = Array(n).fill(false);
  const g: number[][] = Array.from({ length: n }, () => []);
  for (const [a, b] of edges) {
    g[a].push(b);
    g[b].push(a);
  }
  for (const i of restricted) {
    vis[i] = true;
  }
  const q: number[] = [0];
  let ans = 0;
  for (vis[0] = true; q.length; ++ans) {
    const i = q.pop()!;
    for (const j of g[i]) {
      if (!vis[j]) {
        vis[j] = true;
        q.push(j);
      }
    }
  }
  return ans;
}