LeetCode #3715 — HARD

Sum of Perfect Square Ancestors

Break down a hard problem into reliable checkpoints, edge-case handling, and complexity trade-offs.

The Problem

Problem Statement

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] = [u_i, v_i] indicates an undirected edge between nodes u_i and v_i.

You are also given an integer array nums, where nums[i] is the positive integer assigned to node i.

Define a value t_i 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 t_i values for all nodes i in range [1, n - 1].

Note:

In a rooted tree, the ancestors of node i are all nodes on the path from node i to the root node 0, excluding i itself.

Example 1:

Input: n = 3, edges = [[0,1],[1,2]], nums = [2,8,2]

Output: 3

Explanation:

`i`	Ancestors	`nums[i] * nums[ancestor]`	Square Check	`t_i`
1	[0]	`nums[1] * nums[0] = 8 * 2 = 16`	16 is a perfect square	1
2	[1, 0]	`nums[2] * nums[1] = 2 * 8 = 16` `nums[2] * nums[0] = 2 * 2 = 4`	Both 4 and 16 are perfect squares	2

Thus, the total number of valid ancestor pairs across all non-root nodes is 1 + 2 = 3.

Example 2:

Input: n = 3, edges = [[0,1],[0,2]], nums = [1,2,4]

Output: 1

Explanation:

`i`	Ancestors	`nums[i] * nums[ancestor]`	Square Check	`t_i`
1	[0]	`nums[1] * nums[0] = 2 * 1 = 2`	2 is not a perfect square	0
2	[0]	`nums[2] * nums[0] = 4 * 1 = 4`	4 is a perfect square	1

Thus, the total number of valid ancestor pairs across all non-root nodes is 1.

Example 3:

Input: n = 4, edges = [[0,1],[0,2],[1,3]], nums = [1,2,9,4]

Output: 2

Explanation:

`i`	Ancestors	`nums[i] * nums[ancestor]`	Square Check	`t_i`
1	[0]	`nums[1] * nums[0] = 2 * 1 = 2`	2 is not a perfect square	0
2	[0]	`nums[2] * nums[0] = 9 * 1 = 9`	9 is a perfect square	1
3	[1, 0]	`nums[3] * nums[1] = 4 * 2 = 8` `nums[3] * nums[0] = 4 * 1 = 4`	Only 4 is a perfect square	1

Thus, the total number of valid ancestor pairs across all non-root nodes is 0 + 1 + 1 = 2.

Constraints:

1 <= n <= 10⁵
edges.length == n - 1
edges[i] = [u_i, v_i]
0 <= u_i, v_i <= n - 1
nums.length == n
1 <= nums[i] <= 10⁵
The input is generated such that edges represents a valid tree.

Step 01

Brute Force Baseline

Problem summary: 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 between nodes ui and vi. 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]. Note: In a rooted tree, the ancestors of node i are all nodes on the path from node i to the root node 0, excluding i itself.

Baseline thinking

Start with the most direct exhaustive search. That gives a correctness anchor before optimizing.

Pattern signal: Array · Hash Map · Math · Tree

Example 1

3
[[0,1],[1,2]]
[2,8,2]

Example 2

3
[[0,1],[0,2]]
[1,2,4]

Example 3

4
[[0,1],[0,2],[1,3]]
[1,2,9,4]

Step 02

Core Insight

What unlocks the optimal approach

Notice that the product <code>nums[i] * nums[ancestor]</code> is a perfect square if and only if both numbers have the same "square-free kernel" (i.e., after removing all even powers of primes, the remaining product is identical).
Precompute the square-free representation of every node's value using prime factorization up to <code>max(nums[i])</code>.
Perform a DFS from the root. While traversing down the tree, maintain a frequency map of the square-free values of the ancestors.
For each node, the number of valid ancestors equals the count of ancestors with the same square-free value.
Carefully backtrack the frequency map after finishing a subtree to maintain correctness.

Interview move: turn each hint into an invariant you can check after every iteration/recursion step.

Step 03

Algorithm Walkthrough

Iteration Checklist

Define state (indices, window, stack, map, DP cell, or recursion frame).
Apply one transition step and update the invariant.
Record answer candidate when condition is met.
Continue until all input is consumed.

Use the first example testcase as your mental trace to verify each transition.

Step 04

Edge Cases

Minimum Input

Single element / shortest valid input

Validate boundary behavior before entering the main loop or recursion.

Duplicates & Repeats

Repeated values / repeated states

Decide whether duplicates should be merged, skipped, or counted explicitly.

Extreme Constraints

Largest constraint values

Re-check complexity target against constraints to avoid time-limit issues.

Invalid / Corner Shape

Empty collections, zeros, or disconnected structures

Handle special-case structure before the core algorithm path.

Step 05

Full Annotated Code

Source-backed implementations are provided below for direct study and interview prep.

// Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
// Auto-generated Java example from go.
class Solution {
    public void exampleSolution() {
    }
}
// Reference (go):
// // Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
// package main
// 
// // https://space.bilibili.com/206214
// const mx = 100_001
// 
// var core = [mx]int{}
// 
// func init() {
// 	for i := 1; i < mx; i++ {
// 		if core[i] == 0 {
// 			for j := 1; i*j*j < mx; j++ {
// 				core[i*j*j] = i
// 			}
// 		}
// 	}
// }
// 
// func sumOfAncestors(n int, edges [][]int, nums []int) (ans int64) {
// 	g := make([][]int, n)
// 	for _, e := range edges {
// 		x, y := e[0], e[1]
// 		g[x] = append(g[x], y)
// 		g[y] = append(g[y], x)
// 	}
// 
// 	cnt := map[int]int{}
// 	var dfs func(int, int)
// 	dfs = func(x, fa int) {
// 		c := core[nums[x]]
// 		// 本题 x 的祖先不包括 x 自己
// 		ans += int64(cnt[c])
// 		cnt[c]++
// 		for _, y := range g[x] {
// 			if y != fa {
// 				dfs(y, x)
// 			}
// 		}
// 		cnt[c]-- // 恢复现场
// 	}
// 	dfs(0, -1)
// 	return
// }

// Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
package main

// https://space.bilibili.com/206214
const mx = 100_001

var core = [mx]int{}

func init() {
	for i := 1; i < mx; i++ {
		if core[i] == 0 {
			for j := 1; i*j*j < mx; j++ {
				core[i*j*j] = i
			}
		}
	}
}

func sumOfAncestors(n int, edges [][]int, nums []int) (ans int64) {
	g := make([][]int, n)
	for _, e := range edges {
		x, y := e[0], e[1]
		g[x] = append(g[x], y)
		g[y] = append(g[y], x)
	}

	cnt := map[int]int{}
	var dfs func(int, int)
	dfs = func(x, fa int) {
		c := core[nums[x]]
		// 本题 x 的祖先不包括 x 自己
		ans += int64(cnt[c])
		cnt[c]++
		for _, y := range g[x] {
			if y != fa {
				dfs(y, x)
			}
		}
		cnt[c]-- // 恢复现场
	}
	dfs(0, -1)
	return
}

# Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
# Time:  precompute: O(r)
#        runtime:    O(nlogx)
# Space: O(r + n)

import collections


def linear_sieve_of_eratosthenes(n):  # Time: O(n), Space: O(n)
    primes = []
    spf = [-1]*(n+1)  # the smallest prime factor
    for i in xrange(2, n+1):
        if spf[i] == -1:
            spf[i] = i
            primes.append(i)
        for p in primes:
            if i*p > n or p > spf[i]:
                break
            spf[i*p] = p
    return spf


MAX_NUMS = 10**5
SPF = linear_sieve_of_eratosthenes(MAX_NUMS)  # Time: O(sqrt(r)), Space: O(sqrt(r))

# number theory, iterative dfs, freq table
class Solution(object):
    def sumOfAncestors(self, n, edges, nums):
        """
        :type n: int
        :type edges: List[List[int]]
        :type nums: List[int]
        :rtype: int
        """
        def prime_factors(x):
            result = 1
            while x != 1:
                if result%SPF[x] == 0:
                    result //= SPF[x]
                else:
                    result *= SPF[x]
                x //= SPF[x]
            return result

        def iter_dfs():
            result = 0
            stk = [(1, (0, -1))]
            while stk:
                step, args = stk.pop()
                if step == 1:
                    u, p = args
                    x = prime_factors(nums[u])
                    result += cnt[x]
                    cnt[x] += 1
                    stk.append((3, (x,)))
                    stk.append((2, (u, p, 0)))
                elif step == 2:
                    u, p, i = args
                    if i == len(adj[u]):
                        continue
                    stk.append((2, (u, p, i+1)))
                    v = adj[u][i]
                    if v == p:
                        continue
                    stk.append((1, (v, u)))
                elif step == 3:
                    x = args[0]
                    cnt[x] -= 1
            return result

        adj = [[] for _ in xrange(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        cnt = collections.defaultdict(int)
        return iter_dfs()


# Time:  precompute: O(r)
#        runtime:    O(nlogx)
# Space: O(r + n)
import collections


# number theory, dfs, freq table
class Solution2(object):
    def sumOfAncestors(self, n, edges, nums):
        """
        :type n: int
        :type edges: List[List[int]]
        :type nums: List[int]
        :rtype: int
        """
        def prime_factors(x):
            result = 1
            while x != 1:
                if result%SPF[x] == 0:
                    result //= SPF[x]
                else:
                    result *= SPF[x]
                x //= SPF[x]
            return result

        def dfs(u, p):
            x = prime_factors(nums[u])
            result = cnt[x]
            cnt[x] += 1
            for v in adj[u]:
                if v == p:
                    continue
                result += dfs(v, u)
            cnt[x] -= 1
            return result

        adj = [[] for _ in xrange(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        cnt = collections.defaultdict(int)
        return dfs(0, -1)

// Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
// Rust example auto-generated from go reference.
// Replace the signature and local types with the exact LeetCode harness for this problem.
impl Solution {
    pub fn rust_example() {
        // Port the logic from the reference block below.
    }
}

// Reference (go):
// // Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
// package main
// 
// // https://space.bilibili.com/206214
// const mx = 100_001
// 
// var core = [mx]int{}
// 
// func init() {
// 	for i := 1; i < mx; i++ {
// 		if core[i] == 0 {
// 			for j := 1; i*j*j < mx; j++ {
// 				core[i*j*j] = i
// 			}
// 		}
// 	}
// }
// 
// func sumOfAncestors(n int, edges [][]int, nums []int) (ans int64) {
// 	g := make([][]int, n)
// 	for _, e := range edges {
// 		x, y := e[0], e[1]
// 		g[x] = append(g[x], y)
// 		g[y] = append(g[y], x)
// 	}
// 
// 	cnt := map[int]int{}
// 	var dfs func(int, int)
// 	dfs = func(x, fa int) {
// 		c := core[nums[x]]
// 		// 本题 x 的祖先不包括 x 自己
// 		ans += int64(cnt[c])
// 		cnt[c]++
// 		for _, y := range g[x] {
// 			if y != fa {
// 				dfs(y, x)
// 			}
// 		}
// 		cnt[c]-- // 恢复现场
// 	}
// 	dfs(0, -1)
// 	return
// }

// Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
// Auto-generated TypeScript example from go.
function exampleSolution(): void {
}
// Reference (go):
// // Accepted solution for LeetCode #3715: Sum of Perfect Square Ancestors
// package main
// 
// // https://space.bilibili.com/206214
// const mx = 100_001
// 
// var core = [mx]int{}
// 
// func init() {
// 	for i := 1; i < mx; i++ {
// 		if core[i] == 0 {
// 			for j := 1; i*j*j < mx; j++ {
// 				core[i*j*j] = i
// 			}
// 		}
// 	}
// }
// 
// func sumOfAncestors(n int, edges [][]int, nums []int) (ans int64) {
// 	g := make([][]int, n)
// 	for _, e := range edges {
// 		x, y := e[0], e[1]
// 		g[x] = append(g[x], y)
// 		g[y] = append(g[y], x)
// 	}
// 
// 	cnt := map[int]int{}
// 	var dfs func(int, int)
// 	dfs = func(x, fa int) {
// 		c := core[nums[x]]
// 		// 本题 x 的祖先不包括 x 自己
// 		ans += int64(cnt[c])
// 		cnt[c]++
// 		for _, y := range g[x] {
// 			if y != fa {
// 				dfs(y, x)
// 			}
// 		}
// 		cnt[c]-- // 恢复现场
// 	}
// 	dfs(0, -1)
// 	return
// }

Step 06

Interactive Study Demo

Use this to step through a reusable interview workflow for this problem.

Custom checkpoints (one per line)

Press Step or Run All to begin.

Step 07

Complexity Analysis

Time

O(n)

Space

O(h)

Approach Breakdown

LEVEL ORDER

O(n) time

O(n) space

BFS with a queue visits every node exactly once — O(n) time. The queue may hold an entire level of the tree, which for a complete binary tree is up to n/2 nodes = O(n) space. This is optimal in time but costly in space for wide trees.

DFS TRAVERSAL

O(n) time

O(h) space

Every node is visited exactly once, giving O(n) time. Space depends on tree shape: O(h) for recursive DFS (stack depth = height h), or O(w) for BFS (queue width = widest level). For balanced trees h = log n; for skewed trees h = n.

Shortcut: Visit every node once → O(n) time. Recursion depth = tree height → O(h) space.

Coach Notes

Common Mistakes

Review these before coding to avoid predictable interview regressions.

Off-by-one on range boundaries

Wrong move: Loop endpoints miss first/last candidate.

Usually fails on: Fails on minimal arrays and exact-boundary answers.

Fix: Re-derive loops from inclusive/exclusive ranges before coding.

Mutating counts without cleanup

Wrong move: Zero-count keys stay in map and break distinct/count constraints.

Usually fails on: Window/map size checks are consistently off by one.

Fix: Delete keys when count reaches zero.

Overflow in intermediate arithmetic

Wrong move: Temporary multiplications exceed integer bounds.

Usually fails on: Large inputs wrap around unexpectedly.

Fix: Use wider types, modular arithmetic, or rearranged operations.

Forgetting null/base-case handling

Wrong move: Recursive traversal assumes children always exist.

Usually fails on: Leaf nodes throw errors or create wrong depth/path values.

Fix: Handle null/base cases before recursive transitions.