Reachable Nodes With Restrictions
Problem
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.
Examples
Example 1
graph TD classDef reachable fill:#51ff51,stroke:#2b8fb8,stroke-width:2px; classDef restricted fill:#ffe6e6,stroke:#d9534f,stroke-width:1.5px; 0(0):::reachable 1(1):::reachable 2(2):::reachable 3(3):::reachable 4(4):::restricted 5(5):::restricted 6(6) 0 --- 1 1 --- 2 1 --- 3 0 --- 4 0 --- 5 5 --- 6 %% reachable: 0,1,2,3 (4 and 5 are restricted so their branches are blocked)
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
graph TD classDef reachable fill:#51ff51,stroke:#2b8fb8,stroke-width:2px; classDef restricted fill:#ffe6e6,stroke:#d9534f,stroke-width:1.5px; 0(0):::reachable 1(1):::restricted 2(2):::restricted 3(3) 4(4):::restricted 5(5):::reachable 6(6):::reachable 0 --- 1 0 --- 2 0 --- 5 0 --- 4 2 --- 3 5 --- 6 %% reachable: 0,5,6 (1,2,4 restricted)
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 <= 10^5edges.length == n - 1edges[i].length == 20 <= ai, bi < nai != biedgesrepresents a valid tree.1 <= restricted.length < n1 <= restricted[i] < n- All the values of
restrictedare unique.
Solution
Method 1 - BFS from Root Skipping Restricted
Intuition
We are given a tree and a set of restricted nodes. We need to count how many nodes are reachable from node 0 without visiting any restricted node. Since the graph is a tree (acyclic, connected), we can use BFS or DFS to traverse from node 0, skipping restricted nodes.
Approach
- Build an adjacency list for the tree from the edges.
- Use a set for restricted nodes for O(1) lookup.
- Traverse the tree from node 0 using BFS or DFS, skipping any node that is restricted or already visited.
- Count the number of nodes visited.
This approach is efficient because each node and edge is visited at most once.
Code
C++
#include <vector>
#include <unordered_set>
#include <queue>
using namespace std;
class Solution {
public:
int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
vector<vector<int>> adj(n);
for (auto& e : edges) {
adj[e[0]].push_back(e[1]);
adj[e[1]].push_back(e[0]);
}
unordered_set<int> rest(restricted.begin(), restricted.end());
vector<bool> vis(n, false);
queue<int> q;
q.push(0);
vis[0] = true;
int cnt = 0;
while (!q.empty()) {
int u = q.front(); q.pop();
cnt++;
for (int v : adj[u]) {
if (!vis[v] && !rest.count(v)) {
vis[v] = true;
q.push(v);
}
}
}
return cnt;
}
};
Go
package solution
func reachableNodes(n int, edges [][]int, restricted []int) int {
adj := make([][]int, n)
for _, e := range edges {
adj[e[0]] = append(adj[e[0]], e[1])
adj[e[1]] = append(adj[e[1]], e[0])
}
rest := make(map[int]bool)
for _, r := range restricted {
rest[r] = true
}
vis := make([]bool, n)
q := []int{0}
vis[0] = true
cnt := 0
for len(q) > 0 {
u := q[0]
q = q[1:]
cnt++
for _, v := range adj[u] {
if !vis[v] && !rest[v] {
vis[v] = true
q = append(q, v)
}
}
}
return cnt
}
Java
import java.util.*;
class Solution {
public int reachableNodes(int n, int[][] edges, int[] restricted) {
List<List<Integer>> adj = new ArrayList<>();
for (int i = 0; i < n; ++i) adj.add(new ArrayList<>());
for (int[] e : edges) {
adj.get(e[0]).add(e[1]);
adj.get(e[1]).add(e[0]);
}
Set<Integer> rest = new HashSet<>();
for (int r : restricted) rest.add(r);
boolean[] vis = new boolean[n];
Queue<Integer> q = new LinkedList<>();
q.add(0);
vis[0] = true;
int cnt = 0;
while (!q.isEmpty()) {
int u = q.poll();
cnt++;
for (int v : adj.get(u)) {
if (!vis[v] && !rest.contains(v)) {
vis[v] = true;
q.add(v);
}
}
}
return cnt;
}
}
Kotlin
class Solution {
fun reachableNodes(n: Int, edges: Array<IntArray>, restricted: IntArray): Int {
val adj = Array(n) { mutableListOf<Int>() }
for (e in edges) {
adj[e[0]].add(e[1])
adj[e[1]].add(e[0])
}
val rest = restricted.toHashSet()
val vis = BooleanArray(n)
val q = ArrayDeque<Int>()
q.add(0)
vis[0] = true
var cnt = 0
while (q.isNotEmpty()) {
val u = q.removeFirst()
cnt++
for (v in adj[u]) {
if (!vis[v] && v !in rest) {
vis[v] = true
q.add(v)
}
}
}
return cnt
}
}
Python
from collections import deque, defaultdict
class Solution:
def reachableNodes(self, n: int, edges: list[list[int]], restricted: list[int]) -> int:
adj = defaultdict(list)
for a, b in edges:
adj[a].append(b)
adj[b].append(a)
rest = set(restricted)
vis = {0}
q = deque([0])
cnt = 0
while q:
u = q.popleft()
cnt += 1
for v in adj[u]:
if v not in vis and v not in rest:
vis.add(v)
q.append(v)
return cnt
Rust
use std::collections::{VecDeque, HashSet};
impl Solution {
pub fn reachable_nodes(n: i32, edges: Vec<Vec<i32>>, restricted: Vec<i32>) -> i32 {
let n = n as usize;
let mut adj = vec![vec![]; n];
for e in edges.iter() {
let (a, b) = (e[0] as usize, e[1] as usize);
adj[a].push(b);
adj[b].push(a);
}
let rest: HashSet<_> = restricted.iter().map(|&x| x as usize).collect();
let mut vis = vec![false; n];
let mut q = VecDeque::new();
q.push_back(0);
vis[0] = true;
let mut cnt = 0;
while let Some(u) = q.pop_front() {
cnt += 1;
for &v in &adj[u] {
if !vis[v] && !rest.contains(&v) {
vis[v] = true;
q.push_back(v);
}
}
}
cnt
}
}
TypeScript
function reachableNodes(n: number, edges: number[][], restricted: number[]): number {
const adj: number[][] = Array.from({length: n}, () => []);
for (const [a, b] of edges) {
adj[a].push(b);
adj[b].push(a);
}
const rest = new Set(restricted);
const vis = new Array(n).fill(false);
const q: number[] = [0];
vis[0] = true;
let cnt = 0;
while (q.length) {
const u = q.shift()!;
cnt++;
for (const v of adj[u]) {
if (!vis[v] && !rest.has(v)) {
vis[v] = true;
q.push(v);
}
}
}
return cnt;
}
Complexity
- ⏰ Time complexity:
O(n)– Each node and edge is visited at most once. - 🧺 Space complexity:
O(n)– For adjacency list, visited set, and queue.