rust 解法, 执行用时: 52 ms, 内存消耗: 18.9 MB, 提交时间: 2024-03-02 11:39:10

use std::collections::HashSet;

impl Solution {
    pub fn reachable_nodes(n: i32, edges: Vec<Vec<i32>>, restricted: Vec<i32>) -> i32 {
        let r = restricted.into_iter().collect::<HashSet<_>>();
        let mut g = vec![vec![]; n as usize];
        for e in &edges {
            let x = e[0];
            let y = e[1];
            if !r.contains(&x) && !r.contains(&y) {
                g[x as usize].push(y as usize); // 都不受限才连边
                g[y as usize].push(x as usize);
            }
        }
        fn dfs(x: usize, fa: usize, g: &Vec<Vec<usize>>) -> i32 {
            let mut cnt = 1;
            for &y in &g[x] {
                if y != fa {
                    cnt += dfs(y, x, g);
                }
            }
            cnt
        }
        dfs(0, 0, &g)
    }
}

javascript 解法, 执行用时: 320 ms, 内存消耗: 111.2 MB, 提交时间: 2024-03-02 11:38:52

/**
 * @param {number} n
 * @param {number[][]} edges
 * @param {number[]} restricted
 * @return {number}
 */
var reachableNodes = function(n, edges, restricted) {
    const r = new Set(restricted);
    const g = Array(n).fill(null).map(() => []);
    for (const [x, y] of edges) {
        if (!r.has(x) && !r.has(y)) {
            g[x].push(y); // 都不受限才连边
            g[y].push(x);
        }
    }
    let ans = 0;
    function dfs(x, fa) {
        ans++;
        for (const y of g[x]) {
            if (y !== fa) {
                dfs(y, x);
            }
        }
    }
    dfs(0, -1);
    return ans;
};

cpp 解法, 执行用时: 361 ms, 内存消耗: 152.5 MB, 提交时间: 2024-03-02 11:38:35

class Solution {
    vector<vector<int>> g;

    int dfs(int x, int fa) {
        int cnt = 1;
        for (int y : g[x]) {
            if (y != fa) {
                cnt += dfs(y, x);
            }
        }
        return cnt;
    };

public:
    int reachableNodes(int n, vector<vector<int>> &edges, vector<int> &restricted) {
        unordered_set<int> r(restricted.begin(), restricted.end());
        g.resize(n);
        for (auto &e : edges) {
            int x = e[0], y = e[1];
            if (!r.contains(x) && !r.contains(y)) {
                g[x].push_back(y); // 都不受限才连边
                g[y].push_back(x);
            }
        }
        return dfs(0, -1);
    }
};

python3 解法, 执行用时: 292 ms, 内存消耗: 63.6 MB, 提交时间: 2023-06-06 09:29:13

# 定义并查集
class UnionFind:
    def __init__(self, n: int):
        self.n = n
        self.parent = list(range(n))

    def find(self, i: int):
        if self.parent[i] != i:
            self.parent[i] = self.find(self.parent[i])
        return self.parent[i]

    def union(self, i: int, j: int):
        rooti = self.find(i)
        rootj = self.find(j)
        if rooti != rootj:
            if rooti > rootj:
                rooti, rootj = rootj, rooti
                i, j = j, i
            self.parent[rootj] = rooti


class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        uf = UnionFind(n)
        restricted = set(restricted)
        for edge in edges:
            if edge[0] not in restricted and edge[1] not in restricted:
                uf.union(edge[0], edge[1])
        ans = 0
        for i in range(n):
            if uf.find(i) == 0:
                ans += 1
        return ans

java 解法, 执行用时: 51 ms, 内存消耗: 86.7 MB, 提交时间: 2023-06-06 09:27:55

// 将restricted中的点转化为vis数组的已访问状态，建图然后通过BFS从0开始遍历图，统计入队节点数即为答案。
class Solution {
    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        List<Integer>[] adj = new List[n];
        for (int i = 0; i < n; ++i){
            adj[i] = new ArrayList<>();
        }
        // 邻接表建图
        for (int i = 0; i < n - 1; ++i){
            adj[edges[i][0]].add(edges[i][1]);
            adj[edges[i][1]].add(edges[i][0]);
        }
        boolean[] vis = new boolean[n];
        // 处理restricted数组
        for (int num : restricted) vis[num] = true;
        Deque<Integer> q = new LinkedList<>();
        q.addLast(0);
        vis[0] = true;
        int ans = 1;
        while (!q.isEmpty()){
            int cur = q.pollFirst();
            for (int next : adj[cur]){
                if (!vis[next]){
                    q.addLast(next);
                    vis[next] = true;
                    ++ans;
                }
            }
        }
        return ans;
    }
}

golang 解法, 执行用时: 312 ms, 内存消耗: 32.4 MB, 提交时间: 2023-06-06 09:27:17

func reachableNodes(n int, edges [][]int, restricted []int) (ans int) {
	r := make(map[int]bool, len(restricted))
	for _, x := range restricted {
		r[x] = true
	}
	g := make([][]int, n)
	for _, e := range edges {
		x, y := e[0], e[1]
		if !r[x] && !r[y] {
			g[x] = append(g[x], y)
			g[y] = append(g[y], x)
		}
	}
	var dfs func(int, int)
	dfs = func(x, fa int) {
		ans++
		for _, y := range g[x] {
			if y != fa {
				dfs(y, x)
			}
		}
	}
	dfs(0, -1)
	return
}

python3 解法, 执行用时: 384 ms, 内存消耗: 126.2 MB, 提交时间: 2023-06-06 09:26:47

'''
用哈希表记录哪些节点是受限的，建图的时候只有当两个节点都不是受限的才连边。
然后 DFS 这棵树，统计从 0 出发能访问到的节点数，即为答案。
'''
class Solution:
    def reachableNodes(self, n: int, edges: List[List[int]], restricted: List[int]) -> int:
        r = set(restricted)
        g = [[] for _ in range(n)]
        for x, y in edges:
            if x not in r and y not in r:
                g[x].append(y)
                g[y].append(x)
        ans = 0
        def dfs(x: int, fa: int) -> None:
            nonlocal ans
            ans += 1
            for y in g[x]:
                if y != fa:
                    dfs(y, x)
        dfs(0, -1)
        return ans