Problem statement

A directed graph has n nodes numbered 0 to n - 1 and exactly n edges: node i has a single outgoing edge to edges[i] (never to itself).

Starting from a node x, keep following the outgoing edges until you arrive at a node you have already visited in this walk.

Return an array answer where answer[i] is the number of distinct nodes visited by the walk that starts at node i.

Example 1

Input: edges = [1,2,0,0]
Output: [3,3,3,4]
Explanation: Nodes 0, 1, 2 form the cycle 0 → 1 → 2 → 0. The walk from 3 visits 3 and then the whole cycle: 4 nodes.

Example 2

Input: edges = [1,2,3,4,0]
Output: [5,5,5,5,5]
Explanation: All five nodes lie on one cycle.

Example 3

Input: edges = [1,0,0,1,3]
Output: [2,2,3,3,4]
Explanation: The cycle is 0 ↔ 1; nodes 2 and 3 hang one step off it and node 4 two steps.

Constraints

• n == edges.length
• 2 <= n <= 10^5
• 0 <= edges[i] <= n - 1
• edges[i] != i

Hints for Count Visited Nodes in a Directed Graph

1. Hint 1

   Every node has exactly one outgoing edge, so each walk runs down a 'tail' and then loops around exactly one cycle. A node on a cycle visits exactly the cycle's nodes.

2. Hint 2

   Find the nodes that are not on any cycle by peeling: repeatedly remove nodes with no incoming edges (Kahn's algorithm). What survives is exactly the cycles.

3. Hint 3

   For each surviving cycle, every node's answer is the cycle length. Then process the peeled nodes in reverse peeling order: answer[i] = answer[edges[i]] + 1.

How to solve Count Visited Nodes in a Directed Graph

In a functional graph every walk ends in a cycle. Peeling nodes of in-degree zero removes exactly the tail nodes, leaving the cycles. Cycle nodes see their own cycle; a tail node sees itself plus everything its successor sees.

Approach

1. Count in-degrees and queue every node with in-degree 0.
2. Pop a node, record it in the peel order, decrement the in-degree of its successor and queue the successor when it reaches 0.
3. Every node never peeled lies on a cycle: walk each such cycle once to measure its length and assign that length to all its nodes.
4. Go through the peel order backwards and set answer[i] = answer[edges[i]] + 1.

Why it works

A node on a cycle always keeps an incoming edge from its cycle predecessor, so it is never peeled; a tail node eventually loses all incoming edges, because everything pointing into it lies on tails too. Walking from a cycle node visits exactly the cycle; walking from a tail node visits the node itself and then the walk of its successor, which shares no nodes with it. A successor is peeled after its predecessor, so the reverse order computes it first.

Complexity

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

Pitfalls

• There can be several separate cycles — measure each one.
• A recursive DFS along a long tail of 10^5 nodes can overflow the stack; the peel order needs none.
• Count distinct nodes: the walk stops on the first repeat, which is a cycle node already counted.

Reference solution

Python

from typing import List

def countVisitedNodes(edges: List[int]) -> List[int]:
    n = len(edges)
    indeg = [0] * n
    for t in edges:
        indeg[t] += 1
    order = [i for i in range(n) if indeg[i] == 0]
    for u in order:
        v = edges[u]
        indeg[v] -= 1
        if indeg[v] == 0:
            order.append(v)
    ans = [0] * n
    for i in range(n):
        if indeg[i] > 0 and ans[i] == 0:
            cycle = [i]
            x = edges[i]
            while x != i:
                cycle.append(x)
                x = edges[x]
            for c in cycle:
                ans[c] = len(cycle)
    for u in reversed(order):
        ans[u] = ans[edges[u]] + 1
    return ans

Java

public static int[] countVisitedNodes(int[] edges) {
    int n = edges.length;
    int[] indeg = new int[n], ans = new int[n], order = new int[n];
    for (int t : edges) indeg[t]++;
    int tail = 0;
    for (int i = 0; i < n; i++) if (indeg[i] == 0) order[tail++] = i;
    for (int h = 0; h < tail; h++) {
        int v = edges[order[h]];
        if (--indeg[v] == 0) order[tail++] = v;
    }
    for (int s = 0; s < n; s++) {
        if (indeg[s] > 0 && ans[s] == 0) {
            int len = 1;
            for (int x = edges[s]; x != s; x = edges[x]) len++;
            ans[s] = len;
            for (int y = edges[s]; y != s; y = edges[y]) ans[y] = len;
        }
    }
    for (int t = tail - 1; t >= 0; t--) ans[order[t]] = ans[edges[order[t]]] + 1;
    return ans;
}

C++

vector<int> countVisitedNodes(vector<int>& edges) {
    int n = edges.size();
    vector<int> indeg(n, 0), ans(n, 0), order;
    for (int t : edges) indeg[t]++;
    for (int i = 0; i < n; i++) if (indeg[i] == 0) order.push_back(i);
    for (size_t h = 0; h < order.size(); h++) {
        int v = edges[order[h]];
        if (--indeg[v] == 0) order.push_back(v);
    }
    for (int s = 0; s < n; s++) {
        if (indeg[s] > 0 && ans[s] == 0) {
            int len = 1;
            for (int x = edges[s]; x != s; x = edges[x]) len++;
            ans[s] = len;
            for (int y = edges[s]; y != s; y = edges[y]) ans[y] = len;
        }
    }
    for (int t = (int) order.size() - 1; t >= 0; t--) ans[order[t]] = ans[edges[order[t]]] + 1;
    return ans;
}

JavaScript

var countVisitedNodes = function(edges) {
    var n = edges.length;
    var indeg = new Array(n), ans = new Array(n);
    for (var i = 0; i < n; i++) { indeg[i] = 0; ans[i] = 0; }
    for (var j = 0; j < n; j++) indeg[edges[j]]++;
    var order = [];
    for (var k = 0; k < n; k++) if (indeg[k] === 0) order.push(k);
    for (var h = 0; h < order.length; h++) {
        var v = edges[order[h]];
        if (--indeg[v] === 0) order.push(v);
    }
    for (var s = 0; s < n; s++) {
        if (indeg[s] > 0 && ans[s] === 0) {
            var len = 1;
            for (var x = edges[s]; x !== s; x = edges[x]) len++;
            ans[s] = len;
            for (var y = edges[s]; y !== s; y = edges[y]) ans[y] = len;
        }
    }
    for (var t = order.length - 1; t >= 0; t--) ans[order[t]] = ans[edges[order[t]]] + 1;
    return ans;
};

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

More dynamic programming problems

• Maximum Sum BST in Binary Tree Hard
• Binary Tree Cameras Hard
• Binary Tree Maximum Path Sum Hard
• Maximum Product of the Length of Two Palindromic Substrings Hard
• Longest Chunked Palindrome Decomposition Hard
• Time Taken to Mark All Nodes Hard