LeetCode 1766 - Tree of Coprimes

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

Solution

Problem Understanding

This problem gives us a tree with n nodes numbered from 0 to n-1, rooted at node 0. Each node has an associated value from the array nums, and the tree structure is given as a list of edges. For each node i, we are asked to find the closest ancestor whose value is coprime with nums[i]. Two numbers are coprime if their greatest common divisor (gcd) is 1. If no such ancestor exists, the answer should be -1.

The input nums defines the value of each node, while edges defines the tree connections. The expected output is an array ans of size n where ans[i] is the closest coprime ancestor to node i.

The constraints tell us that nums[i] values are small (1 <= nums[i] <= 50) but the tree can be large (n <= 10^5). This suggests we must avoid brute-force ancestor searches for each node, as that could lead to O(n^2) time complexity, which is not feasible for large trees. Important edge cases include nodes with no coprime ancestors (result -1) and nodes where multiple ancestors are coprime (we must pick the closest).

Approaches

The brute-force approach is straightforward: for each node i, traverse its ancestors up to the root, compute gcd(nums[i], nums[ancestor]) for each, and keep track of the closest coprime ancestor. This is correct but inefficient. In the worst case, each node requires traversing O(n) ancestors, giving O(n^2) total time complexity, which is too slow for n = 10^5.

The optimal approach relies on the key insight that nums[i] values are small (≤ 50). We can precompute all coprime pairs from 1 to 50. Then, using depth-first search (DFS), we maintain a mapping from values to the most recent node index along the path from root to the current node. At each node, we check which of the ancestor nodes along the path are coprime with nums[i] and select the one closest to the current node. By backtracking after visiting children, we ensure the mapping always reflects the current DFS path.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n^2)	O(n)	Check each node against all ancestors; too slow for large n
Optimal	O(n * V^2)	O(n + V^2)	Precompute coprime pairs; DFS with path mapping; V = 50

Algorithm Walkthrough

Precompute all coprime pairs for numbers 1 to 50. Store for each number the list of values it is coprime with. This allows constant-time lookup for DFS.
Build the adjacency list representation of the tree from the given edges.
Initialize a mapping last_occurrence from node values to the latest node index along the DFS path.
Initialize the result array ans with -1.
Define a recursive DFS function with parameters node and parent. For each node:

a. Initialize variables to track the closest coprime ancestor: closest_node = -1 and closest_depth = -1.

b. For each value v that is coprime with nums[node], check if v exists in last_occurrence. If it does, update closest_node if its depth is closer than the current closest_depth.

c. Update ans[node] = closest_node.

d. Temporarily record the current node in last_occurrence[nums[node]].

e. Recur DFS on all children that are not the parent.

f. Backtrack by removing the current node from last_occurrence[nums[node]]. 6. Start DFS from root node 0 with parent = -1. 7. Return the ans array.

Why it works: The DFS ensures that last_occurrence always reflects the path from the root to the current node, so when checking coprime ancestors, we are only considering actual ancestors. The precomputed coprime mapping and backtracking maintain correctness and efficiency.

Python Solution

from math import gcd
from typing import List, Dict

class Solution:
    def getCoprimes(self, nums: List[int], edges: List[List[int]]) -> List[int]:
        n = len(nums)
        
        # Precompute coprime relationships
        coprime_map: Dict[int, List[int]] = {i: [] for i in range(1, 51)}
        for i in range(1, 51):
            for j in range(1, 51):
                if gcd(i, j) == 1:
                    coprime_map[i].append(j)
        
        # Build tree adjacency list
        tree = [[] for _ in range(n)]
        for u, v in edges:
            tree[u].append(v)
            tree[v].append(u)
        
        ans = [-1] * n
        last_occurrence: Dict[int, List[int]] = {i: [] for i in range(1, 51)}
        
        def dfs(node: int, parent: int):
            closest_node = -1
            closest_depth = -1
            
            # Find closest coprime ancestor
            for val in coprime_map[nums[node]]:
                if last_occurrence[val]:
                    candidate = last_occurrence[val][-1]
                    if candidate > closest_depth:
                        closest_depth = candidate
                        closest_node = candidate
            
            ans[node] = closest_node
            
            # Record current node in last_occurrence
            last_occurrence[nums[node]].append(node)
            
            for child in tree[node]:
                if child != parent:
                    dfs(child, node)
            
            # Backtrack
            last_occurrence[nums[node]].pop()
        
        dfs(0, -1)
        return ans

The Python solution precomputes coprime pairs for values 1 to 50. The DFS maintains a mapping of the last occurrence of each value along the path. For each node, it finds the closest ancestor with a coprime value and backtracks correctly to keep the path accurate.

Go Solution

package main

import "math"

func gcd(a, b int) int {
	for b != 0 {
		a, b = b, a%b
	}
	return a
}

func getCoprimes(nums []int, edges [][]int) []int {
	n := len(nums)
	ans := make([]int, n)
	for i := range ans {
		ans[i] = -1
	}

	// Precompute coprime map
	coprimeMap := make([][]int, 51)
	for i := 1; i <= 50; i++ {
		for j := 1; j <= 50; j++ {
			if gcd(i, j) == 1 {
				coprimeMap[i] = append(coprimeMap[i], j)
			}
		}
	}

	// Build tree
	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)
	}

	lastOccurrence := make([][]int, 51)

	var dfs func(node, parent int)
	dfs = func(node, parent int) {
		closestNode := -1
		for _, val := range coprimeMap[nums[node]] {
			if len(lastOccurrence[val]) > 0 {
				candidate := lastOccurrence[val][len(lastOccurrence[val])-1]
				if candidate > closestNode {
					closestNode = candidate
				}
			}
		}
		ans[node] = closestNode

		lastOccurrence[nums[node]] = append(lastOccurrence[nums[node]], node)
		for _, child := range tree[node] {
			if child != parent {
				dfs(child, node)
			}
		}
		lastOccurrence[nums[node]] = lastOccurrence[nums[node]][:len(lastOccurrence[nums[node]])-1]
	}

	dfs(0, -1)
	return ans
}

The Go solution mirrors the Python logic. The main differences are slice management for lastOccurrence and explicit loop indexing instead of dictionary lookups.

Worked Examples

Example 1:

nums = [2,3,3,2], edges = [[0,1],[1,2],[1,3]]

Precompute coprime pairs.
DFS starts at node 0. No ancestors exist, ans[0] = -1.
Move to node 1, ancestors: [0]. nums[1] = 3, nums[0] = 2 are coprime, ans[1] = 0.
Node 2, ancestors: [0,1]. nums[2] = 3. Check ancestors: 1 → gcd(3,3)=3 (not coprime), 0 → gcd(3,2)=1 (coprime). Closest ancestor is 0 → ans[2]=0.
Node 3, ancestors [0,1]. nums[3]=2. Check ancestors: 1 → gcd(2,3)=1 (coprime), closest ancestor is