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

题目描述

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

给你一个二维整数数组 edges ，长度为 n - 1 ，其中 edges[i] = [a_i, b_i] 表示树中节点 a_i 和 b_i 之间存在一条边。另给你一个整数数组 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 <= 10⁵
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
a_i != b_i
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\) 为节点数。

Python3JavaC++GoTypeScript

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\) 中。

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