Count visited nodes in a directed graph

Dynamic Programming Graph Theory Memoization

30.3% acceptance

Feb 25, 2026

357

There is a directed graph consisting of n nodes numbered from 0 to n - 1 and n directed edges. You are given a 0-indexed array edges where edges[i] indicates that there is an edge from node i to node edges[i]. Consider the following process on the graph: You start from a node x and keep visiting other nodes through edges until you reach a node that you have already visited before on this same process. Return an array answer where answer[i] is the number of different nodes that you will visit if you perform the process starting from node i.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn count_visited_nodes(edges: Vec<i32>) -> Vec<i32> {
    let n = edges.len();
    let mut ans = vec![0i32; n];
    // State: 0=unvisited, 1=in current path, 2=done
    let mut state = vec![0u8; n];
    let mut cycle_len = vec![0i32; n]; // >0 means in cycle with this length

    for start in 0..n {
      if state[start] == 2 { continue; }
      // Walk until we hit visited node
      let mut path = vec![];
      let mut node = start;
      while state[node] == 0 {
        state[node] = 1;
        path.push(node);
        node = edges[node] as usize;
      }
      if state[node] == 1 {
        // node is in current path - find cycle
        let cycle_start_idx = path.iter().position(|&x| x == node).unwrap();
        let clen = (path.len() - cycle_start_idx) as i32;
        // Mark cycle nodes
        for &u in &path[cycle_start_idx..] {
          cycle_len[u] = clen;
          ans[u] = clen;
          state[u] = 2;
        }
        // Mark tail nodes (before cycle)
        for i in (0..cycle_start_idx).rev() {
          let u = path[i];
          ans[u] = ans[edges[u] as usize] + 1;
          state[u] = 2;
        }
      } else {
        // node is already done
        // Fill tail backwards
        for i in (0..path.len()).rev() {
          let u = path[i];
          ans[u] = ans[edges[u] as usize] + 1;
          state[u] = 2;
        }
      }
    }
    ans
  }
}