Reachable Nodes With Restrictions — Medium Problem & Solution

An undirected tree has n nodes labelled 0 … n - 1 and n - 1 edges. Some nodes are listed in restricted and may not be visited.

Difficulty: Medium
Topics: Arrays, Breadth-First Search, Depth-First Search, Graph, Union Find, Trees
Asked at: Amazon, Google, Accenture
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

An undirected tree has n nodes labelled 0 … n - 1 and n - 1 edges. Some nodes are listed in restricted and may not be visited.

Return how many nodes can be reached from node 0, counting node 0 itself. Node 0 is never restricted.

Example 1

Input: n = 7, edges = [[0,1],[1,2],[3,1],[4,0],[0,5],[5,6]], restricted = [4,5]
Output: 4
Explanation: Nodes 0, 1, 2 and 3; the branches through 4 and 5 are cut off.

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: Only nodes 0, 5 and 6 remain reachable.

Example 3

Input: n = 2, edges = [[0,1]], restricted = [1]
Output: 1

Constraints

2 <= n <= 10^5
edges.length == n - 1
edges[i].length == 2
0 <= edges[i][0], edges[i][1] < n
edges represents a valid tree
1 <= restricted.length < n
1 <= restricted[i] < n
All the values of restricted are unique.

How to solve Reachable Nodes With Restrictions

Mark the restricted nodes, then run a depth-first or breadth-first search from node 0 that skips them, and count what it visits.

Approach

Fill a blocked array from restricted.
Build the adjacency list from the edges, both directions.
Traverse from node 0, never entering a blocked node.
Return the number of nodes visited.

Why it works

In a tree there is exactly one path between any two nodes, so a restricted node on that path makes the far end unreachable — the traversal needs no special handling for cut branches, it simply never walks into one. An iterative stack is used rather than recursion because n reaches 100 000 and a path-shaped tree would blow the call stack.

Complexity

Time — O(n)
Space — O(n)

Pitfalls

A restricted list searched linearly per step is O(n · restricted); mark it in an array instead.
Edges are undirected, so both directions go into the adjacency list.
Node 0 counts towards the answer.

Reference solution

Python

from typing import List

def reachableNodes(n: int, edges: List[List[int]], restricted: List[int]) -> int:
    blocked = [False] * n
    for r in restricted:
        blocked[r] = True
    adj = [[] for _ in range(n)]
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
    if blocked[0]:
        return 0
    seen = [False] * n
    seen[0] = True
    stack = [0]
    total = 0
    while stack:
        u = stack.pop()
        total += 1
        for v in adj[u]:
            if seen[v] or blocked[v]:
                continue
            seen[v] = True
            stack.append(v)
    return total

JavaScript

var reachableNodes = function(n, edges, restricted) {
    var i;
    var blocked = [], adj = [];
    for (i = 0; i < n; i++) { blocked.push(false); adj.push([]); }
    for (i = 0; i < restricted.length; i++) blocked[restricted[i]] = true;
    for (i = 0; i < edges.length; i++) {
        adj[edges[i][0]].push(edges[i][1]);
        adj[edges[i][1]].push(edges[i][0]);
    }
    if (blocked[0]) return 0;
    var seen = [];
    for (i = 0; i < n; i++) seen.push(false);
    seen[0] = true;
    var stack = [0];
    var total = 0;
    while (stack.length > 0) {
        var u = stack.pop();
        total++;
        for (i = 0; i < adj[u].length; i++) {
            var v = adj[u][i];
            if (seen[v] || blocked[v]) continue;
            seen[v] = true;
            stack.push(v);
        }
    }
    return total;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 667 arrays problems · the whole catalogue