Count valid paths in a tree

Math Dynamic Programming Tree Depth-First Search Number Theory

36.0% acceptance

Feb 25, 2026

284

There is an undirected tree with n nodes labeled from 1 to n. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [ui, vi] indicates that there is an edge between nodes ui and vi in the tree. Return the number of valid paths in the tree. A path (a, b) is valid if there exists exactly one prime number among the node labels in the path from a to b. Note that: The path (a, b) is a sequence of distinct nodes starting with node a and ending with node b such that every two adjacent nodes in the sequence share an edge in the tree. Path (a, b) and path (b, a) are considered the same and counted only once.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn count_paths(n: i32, edges: Vec<Vec<i32>>) -> i64 {
    let n = n as usize;
    // Sieve primes up to n
    let mut is_prime = vec![true; n + 1];
    is_prime[0] = false;
    if n >= 1 { is_prime[1] = false; }
    let mut i = 2;
    while i * i <= n {
      if is_prime[i] { let mut j = i*i; while j <= n { is_prime[j] = false; j += i; } }
      i += 1;
    }
    // Build adjacency list
    let mut adj = vec![vec![]; n + 1];
    for e in &edges {
      let u = e[0] as usize;
      let v = e[1] as usize;
      adj[u].push(v);
      adj[v].push(u);
    }
    // Tree DP: for each node, count paths with exactly 1 prime passing through it as the "top" node
    // Use DFS. For each node u, maintain count of paths from u downward with 0 primes and 1 prime.
    // When combining subtrees, count valid pairs.
    let mut ans = 0i64;
    // dp[u] = (cnt0, cnt1) where cntk = # paths from u downward with exactly k primes (including u)
    let mut dp = vec![(0i64, 0i64); n + 1];
    let mut parent = vec![0usize; n + 1];
    let mut order = vec![];
    let mut stack = vec![(1usize, 0usize)];
    while let Some((u, p)) = stack.pop() {
      order.push(u);
      parent[u] = p;
      for &v in &adj[u] {
        if v != p { stack.push((v, u)); }
      }
    }
    for &u in order.iter().rev() {
      // Initial: path of just u
      let (mut acc0, mut acc1) = if is_prime[u] { (0i64, 1i64) } else { (1i64, 0i64) };
      for &v in &adj[u] {
        if v == parent[u] { continue; }
        let (b0, b1) = dp[v];
        // Cross paths through u: acc side (includes u) + child side
        // acc[k] * b[j] where k+j == 1
        ans += acc0 * b1 + acc1 * b0;
        // Merge child into accumulator: path from u going through this child
        // u contributes its prime(1 or 0), child path contributes b0/b1
        if is_prime[u] {
          // acc1 += b0 (path u->x: u=1 prime + child 0 primes = 1 prime)
          // acc2 would get b1 (out of range, ignore)
          acc1 += b0;
          // b1 paths would give 2 primes, can't be needed for target=1
        } else {
          acc0 += b0;
          acc1 += b1;
        }
      }
      dp[u] = (acc0, acc1);
    }
    ans
  }
}