leetcode2368. 受限条件下可到达节点的数目-medium

1 题目：受限条件下可到达节点的数目

官方标定难度：中

现有一棵由 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 中的所有值互不相同

2 solution

dfs 搜索没有限制的节点即可

代码

cpp 复制代码

class Solution {
    static int const N = 1e5 + 1;
    vector<int> e[N];
    bitset<N> rest;

    int dfs(int u, int p) {
        int ans = 1;
        for (int v: e[u]) {
            if (v != p && !rest.test(v)) {
                ans += dfs(v, u);
            }
        }
        return ans;
    }

public:
    int reachableNodes(int n, vector<vector<int>> &edges, vector<int> &restricted) {
        for (auto &x: edges) {
            e[x[0]].push_back(x[1]);
            e[x[1]].push_back(x[0]);
        }

        for (int v: restricted) rest.set(v);

        return dfs(0, -1);
    }
};