Sum of perfect square ancestors

Array Hash Table Math Tree Depth-First Search Counting Number Theory

42.5% acceptance

Feb 24, 2026

You are given an integer n and an undirected tree rooted at node 0 with n nodes numbered from 0 to n - 1. This is represented by a 2D array edges of length n - 1, where edges[i] = [ui, vi] indicates an undirected edge. You are also given an integer array nums, where nums[i] is the positive integer assigned to node i. Define a value ti as the number of ancestors of node i such that the product nums[i] * nums[ancestor] is a perfect square. Return the sum of all ti values for all nodes i in range [1, n - 1].

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn sum_of_ancestors(n: i32, edges: Vec<Vec<i32>>, nums: Vec<i32>) -> i64 {
    let n = n as usize;
    // Reduce each number to its square-free part
    fn square_free(mut x: i32) -> i32 {
      let mut res = 1;
      let mut d = 2;
      while d * d <= x {
        let mut cnt = 0;
        while x % d == 0 { x /= d; cnt += 1; }
        if cnt % 2 == 1 { res *= d; }
        d += 1;
      }
      if x > 1 { res *= x; }
      res
    }
    let sf: Vec<i32> = nums.iter().map(|&v| square_free(v)).collect();
    // nums[i]*nums[j] is perfect square iff sf[i] == sf[j]
    // Build adjacency list
    let mut adj = vec![vec![]; n];
    for e in &edges {
      adj[e[0] as usize].push(e[1] as usize);
      adj[e[1] as usize].push(e[0] as usize);
    }
    // DFS from root 0, track ancestor square-free counts
    let mut ans = 0i64;
    let mut ancestor_count: std::collections::HashMap<i32, i32> = std::collections::HashMap::new();
    // Iterative DFS
    let mut stack: Vec<(usize, usize, bool)> = vec![(0, usize::MAX, false)];
    while let Some((node, parent, returning)) = stack.pop() {
      if returning {
        // Remove from ancestors
        let cnt = ancestor_count.entry(sf[node]).or_insert(0);
        *cnt -= 1;
        if *cnt == 0 { ancestor_count.remove(&sf[node]); }
      } else {
        // Count matching ancestors
        ans += *ancestor_count.get(&sf[node]).unwrap_or(&0) as i64;
        // Add to ancestors
        *ancestor_count.entry(sf[node]).or_insert(0) += 1;
        // Push returning marker
        stack.push((node, parent, true));
        // Push children
        for &child in &adj[node] {
          if child != parent {
            stack.push((child, node, false));
          }
        }
      }
    }
    ans
  }
}