【哈希表, 图, 并查集, DFS】受限条件下可到达节点的数目

题目

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

---

现有一棵由 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 <= 10^5
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• edges 表示一棵有效的树
• 1 <= restricted.length < n
• 1 <= restricted[i] < n
• restricted 中的所有值 互不相同

解题

方法一：哈希表 图 并查集

思路

把该无向二叉树看作一个无向图，那么以这个无向图维护成一个并查集，找出有多少个节点与节点 0 相连就是答案。

受限节点数组(restricted)维护成一个哈希表，便于判断某个节点是否为受限节点（这里用一个长度为 n 的布尔数组替代）。遍历边数组(edges)，当两节点都不是受限节点时把这对节点合并。最后遍历所有可能节点，统计出与节点 0 相连的节点数量返回。

代码

class Solution {
    private int[] root;

    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        root = new int[n];
        boolean[] forbid = new boolean[n];
        for (int i = 0; i < n; ++i) root[i] = i;
        for (int v : restricted) forbid[v] = true;
        for (int[] edge : edges) {
            if (!forbid[edge[0]] && !forbid[edge[1]]) {
                union(edge[0], edge[1]);
            }
        }
        int count = 0, root = find(0);
        for (int i = 0; i < n; ++i) {
            if (find(i) == root) ++count;
        }
        return count;
    }

    private int find(int n) {
        return root[n] == n ? n : (root[n] = find(root[n]));
    }

    private void union(int p, int q) {
        root[find(p)] = find(q);
    }
}

int _ = []() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    return 0;
}();

class Solution {
    vector<int> root;
    vector<bool> forbid;

    int find(int n) {
        return root[n] == n ? n : (root[n] = find(root[n]));
    }

    void join(int p, int q) {
        root[find(p)] = find(q);
    }

public:
    int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
        root.resize(n), forbid.resize(n);
        for (int i = 0; i < n; ++i) root[i] = i;
        for (int& v : restricted) forbid[v] = true;
        for (vector<int>& edge : edges) {
            if (!forbid[edge[0]] && !forbid[edge[1]]) {
                join(edge[0], edge[1]);
            }
        }
        int count = 0, root = find(0);
        for (int i = 0; i < n; ++i) {
            if (find(i) == root) ++count;
        }
        return count;
    }
};

方法二：DFS

思路

把边数组(edges)维护成一个邻接表(graph)，然后深搜。

代码

class Solution {
    vector<vector<int>> graph;
    vector<bool> forbid;
    int count;

    void dfs(int u, int v) {
        ++count;
        for (int& con : graph[u]) {
            if (con != v && !forbid[con]) dfs(con, u);
        }
    }

public:
    int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
        graph.resize(n), forbid.resize(n);
        for (vector<int>& edge : edges) {
            graph[edge[0]].push_back(edge[1]);
            graph[edge[1]].push_back(edge[0]);
        }
        for (int& v : restricted) forbid[v] = true;
        dfs(0, -1);

        return count;
    }
};