2368. Reachable Nodes With Restrictions

Description

There is an undirected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given a 2D integer array edges of length n - 1 where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree. You are also given an integer array restricted which represents restricted nodes.

Return the maximum number of nodes you can reach from node 0 without visiting a restricted node.

Note that node 0 will not be a restricted node.

 

Example 1:

Input: n = 7, edges = [[0,1],[1,2],[3,1],[4,0],[0,5],[5,6]], restricted = [4,5]
Output: 4
Explanation: The diagram above shows the tree.
We have that [0,1,2,3] are the only nodes that can be reached from node 0 without visiting a restricted node.

Example 2:

Input: n = 7, edges = [[0,1],[0,2],[0,5],[0,4],[3,2],[6,5]], restricted = [4,2,1]
Output: 3
Explanation: The diagram above shows the tree.
We have that [0,5,6] are the only nodes that can be reached from node 0 without visiting a restricted node.

 

Constraints:

  • 2 <= n <= 105
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • ai != bi
  • edges represents a valid tree.
  • 1 <= restricted.length < n
  • 1 <= restricted[i] < n
  • All the values of restricted are unique.

Solutions

Solution 1

Python Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution:
    def reachableNodes(
        self, n: int, edges: List[List[int]], restricted: List[int]
    ) -> int:
        g = defaultdict(list)
        vis = [False] * n
        for v in restricted:
            vis[v] = True
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)

        def dfs(u):
            nonlocal ans
            if vis[u]:
                return
            ans += 1
            vis[u] = True
            for v in g[u]:
                dfs(v)

        ans = 0
        dfs(0)
        return ans

Java Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
    private List<Integer>[] g;
    private boolean[] vis;
    private int ans;

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

        ans = 0;
        dfs(0);
        return ans;
    }

    private void dfs(int u) {
        if (vis[u]) {
            return;
        }
        ++ans;
        vis[u] = true;
        for (int v : g[u]) {
            dfs(v);
        }
    }
}

C++ Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
public:
    int ans;

    int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        vector<bool> vis(n);
        for (int v : restricted) vis[v] = true;
        ans = 0;
        dfs(0, g, vis);
        return ans;
    }

    void dfs(int u, vector<vector<int>>& g, vector<bool>& vis) {
        if (vis[u]) return;
        vis[u] = true;
        ++ans;
        for (int v : g[u]) dfs(v, g, vis);
    }
};

Go Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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
}

TypeScript Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
function reachableNodes(n: number, edges: number[][], restricted: number[]): number {
    let res = 0;
    const vis = new Array(n).fill(false);
    const map = new Map<number, number[]>();
    for (const [start, end] of edges) {
        map.set(start, [...(map.get(start) ?? []), end]);
        map.set(end, [...(map.get(end) ?? []), start]);
    }
    const dfs = (cur: number) => {
        if (restricted.includes(cur) || vis[cur]) {
            return;
        }
        res++;
        vis[cur] = true;
        for (const item of map.get(cur) ?? []) {
            dfs(item);
        }
    };
    dfs(0);

    return res;
}

Solution 2

Python Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
    def reachableNodes(
        self, n: int, edges: List[List[int]], restricted: List[int]
    ) -> int:
        s = set(restricted)
        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        q = deque([0])
        vis = [False] * n
        for v in restricted:
            vis[v] = True
        ans = 0
        while q:
            i = q.popleft()
            ans += 1
            vis[i] = True
            for j in g[i]:
                if not vis[j]:
                    q.append(j)
        return ans

Java Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Solution {
    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        List<Integer>[] g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (int[] e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        boolean[] vis = new boolean[n];
        for (int v : restricted) {
            vis[v] = true;
        }
        Deque<Integer> q = new ArrayDeque<>();
        q.offer(0);
        int ans = 0;
        while (!q.isEmpty()) {
            int i = q.pollFirst();
            ++ans;
            vis[i] = true;
            for (int j : g[i]) {
                if (!vis[j]) {
                    q.offer(j);
                }
            }
        }
        return ans;
    }
}

C++ Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
public:
    int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
        vector<vector<int>> g(n);
        vector<bool> vis(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        for (int v : restricted) vis[v] = true;
        queue<int> q{{0}};
        int ans = 0;
        while (!q.empty()) {
            int i = q.front();
            q.pop();
            ++ans;
            vis[i] = true;
            for (int j : g[i])
                if (!vis[j]) q.push(j);
        }
        return ans;
    }
};

Go Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
func reachableNodes(n int, edges [][]int, restricted []int) 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 _, v := range restricted {
		vis[v] = true
	}
	q := []int{0}
	ans := 0
	for len(q) > 0 {
		i := q[0]
		q = q[1:]
		ans++
		vis[i] = true
		for _, j := range g[i] {
			if !vis[j] {
				q = append(q, j)
			}
		}
	}
	return ans
}

TypeScript Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
function reachableNodes(n: number, edges: number[][], restricted: number[]): number {
    const g = Array.from({ length: n }, () => []);
    const vis = new Array(n).fill(false);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    for (const v of restricted) {
        vis[v] = true;
    }
    const q = [0];
    let ans = 0;
    while (q.length) {
        const i = q.shift();
        ++ans;
        vis[i] = true;
        for (const j of g[i]) {
            if (!vis[j]) {
                q.push(j);
            }
        }
    }
    return ans;
}