Reachable Nodes With Restrictions

medium graph tree bfs dfs

Problem

You are given an undirected tree with n nodes labeled 0..n-1 described by n-1 edges, and an array restricted of forbidden nodes. Starting at node 0 (which is never restricted), return the maximum number of nodes you can reach without ever visiting a restricted node.

Inputn = 7, edges = [[0,1],[1,2],[3,1],[4,0],[0,5],[5,6]], restricted = [4,5]

Output4

Nodes 4 and 5 are blocked, so 5 and 6 are cut off. Only [0, 1, 2, 3] are reachable.

Inputn = 7, edges = [[0,1],[0,2],[0,5],[0,4],[3,2],[6,5]], restricted = [4,2,1]

Output3

Only [0, 5, 6] survive; every other branch passes through a blocked node.

n (number of nodes)

edges (a:b, comma-separated)

restricted (comma-separated)

def reachable_nodes(n, edges, restricted):
    blocked = set(restricted)
    adj = [[] for _ in range(n)]
    for a, b in edges:
        adj[a].append(b)
        adj[b].append(a)
    seen = {0}
    queue = [0]
    count = 0
    while queue:
        u = queue.pop()
        count += 1
        for v in adj[u]:
            if v not in seen and v not in blocked:
                seen.add(v)
                queue.append(v)
    return count

function reachableNodes(n, edges, restricted) {
  const blocked = new Set(restricted);
  const adj = Array.from({ length: n }, () => []);
  for (const [a, b] of edges) {
    adj[a].push(b);
    adj[b].push(a);
  }
  const seen = new Set([0]);
  const queue = [0];
  let count = 0;
  while (queue.length) {
    const u = queue.pop();
    count++;
    for (const v of adj[u]) {
      if (!seen.has(v) && !blocked.has(v)) {
        seen.add(v);
        queue.push(v);
      }
    }
  }
  return count;
}

int reachableNodes(int n, int[][] edges, int[] restricted) {
    Set<Integer> blocked = new HashSet<>();
    for (int r : restricted) blocked.add(r);
    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]);
    }
    boolean[] seen = new boolean[n];
    seen[0] = true;
    Deque<Integer> stack = new ArrayDeque<>();
    stack.push(0);
    int count = 0;
    while (!stack.isEmpty()) {
        int u = stack.pop();
        count++;
        for (int v : adj.get(u)) {
            if (!seen[v] && !blocked.contains(v)) {
                seen[v] = true;
                stack.push(v);
            }
        }
    }
    return count;
}

int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
    unordered_set<int> blocked(restricted.begin(), restricted.end());
    vector<vector<int>> adj(n);
    for (auto& e : edges) {
        adj[e[0]].push_back(e[1]);
        adj[e[1]].push_back(e[0]);
    }
    vector<bool> seen(n, false);
    seen[0] = true;
    vector<int> stack = {0};
    int count = 0;
    while (!stack.empty()) {
        int u = stack.back(); stack.pop_back();
        count++;
        for (int v : adj[u]) {
            if (!seen[v] && !blocked.count(v)) {
                seen[v] = true;
                stack.push_back(v);
            }
        }
    }
    return count;
}

Explanation

The graph is a tree, so there is exactly one path between any two nodes. The set of nodes reachable from node 0 is simply the connected component containing 0 after we delete every restricted node. We never need to worry about cycles or multiple routes — a single traversal counts the component.

First, drop the restricted labels into a hash set blocked for O(1) membership tests, and build an adjacency list adj from the undirected edges (each edge is added in both directions).

Then run a graph traversal from node 0. We use a stack here (iterative DFS), but a queue (BFS) gives the same count — order does not matter, only which nodes are visitable. We mark 0 as seen up front so it is counted once.

Pop a node u, increment count, and push each neighbor v that is neither already seen nor blocked. The two guards are the whole trick: skipping blocked nodes prevents us from ever crossing a restricted node, which automatically severs everything behind it.

When the stack empties, count is the number of nodes we managed to enter — the answer. In the default example, nodes 4 and 5 are blocked, so the branch 5 → 6 is unreachable and only [0, 1, 2, 3] are counted, giving 4.

Time: O(n) Space: O(n)