Tree of coprimes

Array Math Tree Depth-First Search Number Theory

43.8% acceptance

Feb 25, 2026

427

There is a tree consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges. Each node has a value associated with it, and the root of the tree is node 0. Two values x and y are coprime if gcd(x, y) == 1. Return an array ans of size n, where ans[i] is the closest ancestor to node i such that nums[i] and nums[ans[i]] are coprime, or -1 if there is no such ancestor.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn get_coprimes(nums: Vec<i32>, edges: Vec<Vec<i32>>) -> Vec<i32> {
    let n = nums.len();
    let mut adj = vec![vec![]; n];
    for e in &edges {
      let (u, v) = (e[0] as usize, e[1] as usize);
      adj[u].push(v);
      adj[v].push(u);
    }

    // Precompute coprimes for values 1..=50
    let mut coprimes = vec![vec![]; 51];
    for v in 1usize..=50 {
      for u in 1usize..=50 {
        if Self::gcd(u as i32, v as i32) == 1 {
          coprimes[v].push(u);
        }
      }
    }

    // For each value 1..=50, stack of (node, depth)
    let mut stacks: Vec<Vec<(i32, i32)>> = vec![vec![]; 51];
    let mut ans = vec![-1i32; n];

    fn dfs(
      node: usize, parent: usize, d: i32,
      adj: &Vec<Vec<usize>>,
      nums: &Vec<i32>,
      coprimes: &Vec<Vec<usize>>,
      stacks: &mut Vec<Vec<(i32, i32)>>,
      ans: &mut Vec<i32>,
    ) {
      let val = nums[node] as usize;
      let mut best_node = -1i32;
      let mut best_depth = -1i32;
      for &co in &coprimes[val] {
        if let Some(&(nd, dd)) = stacks[co].last() {
          if dd > best_depth {
            best_depth = dd;
            best_node = nd;
          }
        }
      }
      ans[node] = best_node;
      stacks[val].push((node as i32, d));
      for &child in &adj[node] {
        if child != parent {
          dfs(child, node, d + 1, adj, nums, coprimes, stacks, ans);
        }
      }
      stacks[val].pop();
    }

    dfs(0, usize::MAX, 0, &adj, &nums, &coprimes, &mut stacks, &mut ans);
    ans
  }

  fn gcd(a: i32, b: i32) -> i32 {
    if b == 0 { a } else { Self::gcd(b, a % b) }
  }
}