LeetCode 3715 - Sum of Perfect Square Ancestors

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Math, Tree, Depth-First Search, Counting, Number Theory

Solution

Problem Understanding

This problem asks us to compute a sum across all non-root nodes in a tree based on ancestor relationships and the property of perfect squares. We are given a rooted tree of n nodes, represented as an edge list. Each node has a positive integer value from the nums array. For each node i (excluding the root), we need to count how many of its ancestors j satisfy the property that nums[i] * nums[j] is a perfect square. Finally, we sum all these counts for nodes 1 through n-1.

The key points are: the tree is rooted at node 0, the tree structure guarantees no cycles, and nums[i] can go up to 10^5. Given n can be up to 10^5, a brute-force solution that checks each node against all ancestors would be too slow.

Edge cases to watch out for include very deep trees (chains), nodes with the value 1, repeated numbers in nums, and large prime numbers, since perfect square conditions depend on prime factorization.

Approaches

Brute-Force Approach

The simplest approach is to traverse the tree for each node, collect all its ancestors, and check whether the product with each ancestor is a perfect square. Checking if a number is a perfect square can be done using integer square root checks. While correct, this is inefficient because a node at depth d requires O(d) checks, leading to worst-case O(n^2) for a chain-like tree. This will not scale for n = 10^5.

Optimal Approach

The key insight is to use prime factorization modulo 2. A product of two numbers is a perfect square if all prime exponents in its factorization sum to even numbers. If we maintain the XOR (mod 2) of prime exponents along the path from root to a node, we can efficiently determine whether the product with any ancestor results in a perfect square. Specifically, we:

Factorize each number into its primes modulo 2 (only keep whether the exponent is odd).
Use a DFS traversal, maintaining a hash map that counts how many times each factorization pattern has been seen on the current path.
For a node i, compute its factorization modulo 2 and check how many ancestors have the same factorization (because XOR of identical vectors yields zero, a perfect square).

This reduces the problem from O(n^2) to O(n log max_num) due to factorization and tree traversal.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n^2)	O(n)	Checks all ancestors explicitly, too slow for n=10^5
Optimal	O(n log max_num)	O(n log max_num)	Uses DFS + prime factorization modulo 2 + hash map

Algorithm Walkthrough

Precompute all primes up to 10^5 using the Sieve of Eratosthenes. This allows efficient factorization.
Factorize each nums[i] into prime exponents modulo 2. Represent it as a tuple of (prime, exponent % 2) or a bitmask.
Construct the adjacency list of the tree from the edges.
Initialize a hash map counter to store counts of each factorization along the current DFS path.
Start DFS from the root node:
For each child, compute the combined factorization with its parent.
Increment ti for the child by the count of nodes along the path with the same factorization pattern (from the hash map).
Update the hash map for this node.
Recursively traverse children.
After recursion, decrement the hash map to backtrack.
Sum all ti values to get the result.

Why it works: By maintaining the XOR of prime exponents along the path and counting matches, we efficiently count ancestor nodes whose product with the current node is a perfect square. Factorization modulo 2 captures the necessary condition for perfect squares.

Python Solution

from typing import List
from collections import defaultdict
import math

class Solution:
    def sumOfAncestors(self, n: int, edges: List[List[int]], nums: List[int]) -> int:
        # Sieve to get primes up to 10^5
        MAX_NUM = 10**5
        sieve = [0] * (MAX_NUM + 1)
        primes = []
        for i in range(2, MAX_NUM + 1):
            if sieve[i] == 0:
                primes.append(i)
                for j in range(i, MAX_NUM + 1, i):
                    if sieve[j] == 0:
                        sieve[j] = i

        def factorize(x: int):
            res = defaultdict(int)
            while x > 1:
                p = sieve[x]
                res[p] ^= 1  # keep only parity
                x //= p
            return tuple(sorted((k,v) for k,v in res.items() if v))

        num_factors = [factorize(num) for num in nums]

        # Build adjacency list
        tree = [[] for _ in range(n)]
        for u, v in edges:
            tree[u].append(v)
            tree[v].append(u)

        ans = 0
        counter = defaultdict(int)

        def dfs(node: int, parent: int, path_factor: tuple):
            nonlocal ans
            # Count ancestors with same factorization
            ans += counter[path_factor]
            # Include current node in path
            counter[num_factors[node]] += 1
            for child in tree[node]:
                if child != parent:
                    combined = tuple(sorted((k, v) for k, v in dict(num_factors[child]).items()))
                    dfs(child, node, num_factors[node])
            counter[num_factors[node]] -= 1

        # Start DFS from root
        counter[factorize(nums[0])] = 1
        for child in tree[0]:
            dfs(child, 0, factorize(nums[0]))

        return ans

Explanation:

We first precompute prime factorizations using a sieve for speed. Each number is factorized into primes modulo 2. Using DFS, we track counts of each factorization pattern along the path, and for each node, we add the count of ancestors with the same pattern. This captures all products that are perfect squares.

Go Solution

package main

func sumOfAncestors(n int, edges [][]int, nums []int) int64 {
    const MAX_NUM = 100000
    sieve := make([]int, MAX_NUM+1)
    primes := []int{}
    for i := 2; i <= MAX_NUM; i++ {
        if sieve[i] == 0 {
            primes = append(primes, i)
            for j := i; j <= MAX_NUM; j += i {
                if sieve[j] == 0 {
                    sieve[j] = i
                }
            }
        }
    }

    factorize := func(x int) map[int]int {
        res := make(map[int]int)
        for x > 1 {
            p := sieve[x]
            res[p] ^= 1
            x /= p
        }
        for k, v := range res {
            if v == 0 {
                delete(res, k)
            }
        }
        return res
    }

    numFactors := make([]map[int]int, n)
    for i := 0; i < n; i++ {
        numFactors[i] = factorize(nums[i])
    }

    tree := make([][]int, n)
    for _, e := range edges {
        u, v := e[0], e[1]
        tree[u] = append(tree[u], v)
        tree[v] = append(tree[v], u)
    }

    counter := map[string]int{}
    var encode func(map[int]int) string
    encode = func(m map[int]int) string {
        s := ""
        for k, v := range m {
            s += string(rune(k)) + ":" + string(rune(v)) + ","
        }
        return s
    }

    var ans int64
    var dfs func(int, int)
    dfs = func(node, parent int) {
        key := encode(numFactors[node])
        ans += int64(counter[key])
        counter[key]++
        for _, child := range tree[node] {
            if child != parent {
                dfs(child, node)
            }
        }
        counter[key]--
    }

    counter[encode(numFactors[0])] = 1
    for _, child := range tree[0] {
        dfs(child, 0)
    }

    return ans
}

Explanation:

Go version mirrors Python. We use a map to store the parity of prime exponents, and string encoding to use maps as keys. DFS tracks counts and sums valid ancestor matches. Integer operations and map management are Go-specific.

Worked Examples

Example 1:

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

Factorization modulo 2:

2 = {2:1}
8 = {2:1} (since 8 = 2^3, parity 3 % 2 = 1)

DFS:

Node 1: ancestor 0 has {2:1}, product parity XOR {2:1} ^ {2:1} = {2:0}