1766. Tree of Coprimes


Problem Description

In this problem, you are presented with a tree structure which is an undirected graph that does not have cycles. The tree is made up of n nodes, each uniquely numbered from 0 to n - 1, and n - 1 edges which connect these nodes. The node with the number 0 is designated as the root of the tree.

Each node is associated with a value, and these values are represented by an integer array nums, where the value of the i-th node is nums[i]. Additionally, the tree structure is defined by a 2D array edges, where each element edges[j] is a two-element array representing an edge connecting the nodes u_j to v_j.

The concept of coprimality is central to the problem. Two values are said to be coprime if their greatest common divisor (gcd) is 1.

An ancestor of a node i is defined as any node that lies in the shortest path from node i back to the root, excluding the node itself.

The objective is to find for each node i, the closest ancestor whose value is coprime with the value nums[i]. The result should be returned as an array of size n, named ans, where ans[i] holds the number of the closest coprime ancestor or -1 if no such ancestor exists.

Intuition

The key to solving this problem is to perform a Depth-First Search (DFS) traversal of the tree, during which we keep track of the ancestors for each node and check their coprimality with the current node's value.

We'll utilize a recursive DFS function that traverses the tree from the root to its leaves, keeping track of each node's ancestors' values that are coprime with the current node's value. To reduce the time complexity, it's crucial to avoid checking coprimality with every ancestor for every node. One way to be efficient is to check only the latest added ancestor that still hasn't been found to be non-coprime, instead of all ancestors.

To implement this, we can maintain a list (stack) for each value v that appears in nums. When we visit a node i with a value of v, we push the node and its depth onto stack v. As we backtrack in the DFS (after visiting all children of a node), we pop the node off stack v.

During the DFS, for each node i, we'll look at the stacks associated with numbers that are known to be coprime with nums[i]. The closest ancestor with a coprime value is the one with the highest depth from among these stacks. We remember this ancestor as the solution for the current node.

To speed up the process of identifying coprime numbers, we precompute coprime pairs for all numbers within the possible range of nums using the gcd function. This way, we can directly reference the precomputed coprime numbers without recalculating the gcd during the DFS.

Overall, the intuition behind the solution is to systematically traverse the tree while efficiently tracking ancestors' values so that we can quickly identify the closest coprime ancestor for each node.

Solution Approach

The solution uses the Depth-First Search (DFS) algorithm to traverse the tree. In the process, it also utilizes some additional data structures and concepts to optimize the search for coprime ancestors.

Data Structures Used:

  • Adjacency List (g): A defaultdict of lists used to store the tree structure in an adjacency list format. This is populated using the edges array and allows efficient traversal of connected nodes.

  • Coprimes Map (f): A defaultdict of lists, precomputed to hold lists of coprime numbers for each possible value in nums. This map is created by iterating through all pairs of numbers within the range [1, 50] and storing pairs that are coprime (where gcd(i, j) == 1) together.

  • Stacks (stks): A defaultdict of lists (acting as stacks) to keep track of the nodes for each unique value found in nums. Each entry in a stack is a tuple of a node and its depth (node, depth). This aids in storing and retrieving the most recently visited ancestor with a particular value.

  • Result List (ans): An array initialized to -1's of length equal to the number of nodes n, used to store the closest coprime ancestor for each node.

Algorithm and Patterns:

  1. The DFS is initialized at the root of the tree (node 0) with an initial depth of 0.

  2. For each node i visited by the DFS, we attempt to find the closest ancestor whose value is coprime with nums[i]. To do this, we look into the stacks for values that are known to be coprime with nums[i] from the f map.

  3. We examine all stacks associated with these coprime values, and among all the nodes in these stacks, we find the one with the maximum depth. This node will be the closest ancestor with the coprime value, and we record it in ans[i].

  4. Before going deeper into the DFS to visit children nodes (j) of the current node i, we push (i, depth) onto the stack associated with nums[i].

  5. After the DFS call returns from visiting a child j, we pop the current node i from the corresponding stack to backtrack.

  6. This process is repeated until all nodes are visited. The stks structure ensures that at each node, we only consider ancestors that are in the path from the current node to the root.

By using the precomputed coprime number map and maintaining the stacks for quick ancestor tracking, the algorithm efficiently searches for the closest coprime ancestor for each node without redundant computations.

Upon completion, the ans array is returned, which contains the closest coprime ancestor for each node or -1 if no such ancestor exists.

Example Walkthrough

Let's consider a small example to illustrate the solution approach. Say we have n = 5 nodes with values nums = [2, 3, 6, 2, 5] in a tree with edges defined as [[0, 1], [1, 2], [1, 3], [2, 4]]. The tree structure would look like this:

1    2 (0)
2    |
3    3 (1)
4   / \
56 (2) 2 (3)
6  |
75 (4)

The numbers in the parentheses represent node numbers.

  1. Construct coprimes map (f) by precomputing coprime pairs for all numbers within the range [1, 50].

    For instance, f[2] = [1, 3, 5, ...], since 2 is coprime with 1, 3, 5, and so on.

  2. Initialize the adjacency list (g) and populate it using the edges array:

    g[0] = [1], g[1] = [0, 2, 3], g[2] = [1, 4], g[3] = [1], g[4] = [2].

  3. Initialize stacks (stks) and result list (ans) with -1 for all nodes.

  4. Perform DFS starting from the root (node 0) with depth = 0.

    • Visit node 0 and push it onto the stack for value 2 as (0, 0).
    • Since node 0 has no ancestors, ans[0] remains -1.
  5. DFS on node 1:

    • Push (1, 1) onto the stack for value 3.
    • Look for stacks of values coprime with 3 using f[3]. We find that 2 is coprime with 3, and stack for value 2 has (0,0). The highest depth is 0, so ans[1] = 0.
  6. DFS on node 2:

    • Push (2, 2) onto the stack for value 6.
    • Look for stacks of values coprime with 6 using f[6]. We find that 5 is coprime with 6, but there is no stack for 5 yet. Check for 3, and we find (1, 1) with depth 1. So, ans[2] = 1.
  7. DFS on node 3:

    • Push (3, 2) onto the stack for value 2.
    • 2 is not coprime with itself, and there are no other ancestors with coprime values, so ans[3] = -1.
  8. Return to node 1 and DFS on node 4:

    • Push (4, 3) onto the stack for value 5.
    • Look for stacks of values coprime with 5 using f[5]. We find that 2 and 3 are coprime with 5. The stack for 2 has a depth of 0 from node 0, and the stack for 3 has a depth of 1 from node 1. So, ans[4] = 1 because node 1 has a higher depth.
  9. Pop nodes from stacks as we backtrack in the DFS call.

Upon completion, the ans array is [-1, 0, 1, -1, 1], indicating the closest coprime ancestors for each node. This result can now be returned.

Python Solution

1from math import gcd
2from collections import defaultdict
3
4class Solution:
5    def getCoprimes(self, nums: List[int], edges: List[List[int]]) -> List[int]:
6        # Function to execute Depth-First Search (DFS)
7        def dfs(current_node, parent_node, depth):
8            closest_coprime_ancestor = -1
9            max_depth = -1
10            for coprime_number in coprimes[nums[current_node]]:
11                stack = stacks[coprime_number]
12                if stack and stack[-1][1] > max_depth:
13                    closest_coprime_ancestor, max_depth = stack[-1]
14            ancestors[current_node] = closest_coprime_ancestor
15            for neighbor in graph[current_node]:
16                if neighbor != parent_node:
17                    stacks[nums[current_node]].append((current_node, depth))
18                    # Recursive DFS call
19                    dfs(neighbor, current_node, depth + 1)
20                    # Pop the element when backtracking
21                    stacks[nums[current_node]].pop()
22
23        # Build the graph from the edges
24        graph = defaultdict(list)
25        for src, dest in edges:
26            graph[src].append(dest)
27            graph[dest].append(src)
28
29        # Precompute coprimes for numbers from 1 to 50
30        coprimes = defaultdict(list)
31        for i in range(1, 51):
32            for j in range(1, 51):
33                if gcd(i, j) == 1:
34                    coprimes[i].append(j)
35
36        # Initialize stacks to keep track of ancestors' values
37        stacks = defaultdict(list)
38        # Initialize the list to store ancestors for each node
39        ancestors = [-1] * len(nums)
40
41        # Start DFS from the root node, here considered to be node 0
42        dfs(0, -1, 0)
43      
44        # Return the final list of ancestors
45        return ancestors
46

Java Solution

1class Solution {
2    // Graph representation using adjacency lists for nodes and their co-prime lists
3    private List<Integer>[] graph;
4    private List<Integer>[] coPrimeLists;
5  
6    // Stacks to maintain the ancestors of each node that satisfy the co-prime condition
7    private Deque<int[]>[] ancestorStacks;
8  
9    // Array of numbers assigned to each node
10    private int[] nums;
11  
12    // Array to store the final answers
13    private int[] answers;
14
15    public int[] getCoprimes(int[] nums, int[][] edges) {
16        int n = nums.length;
17        graph = new List[n];
18        Arrays.setAll(graph, k -> new ArrayList<>());
19      
20        // Building the undirected graph from the edges
21        for (var edge : edges) {
22            int u = edge[0], v = edge[1];
23            graph[u].add(v);
24            graph[v].add(u);
25        }
26      
27        coPrimeLists = new List[51];
28        ancestorStacks = new Deque[51];
29        Arrays.setAll(coPrimeLists, k -> new ArrayList<>());
30        Arrays.setAll(ancestorStacks, k -> new ArrayDeque<>());
31      
32        // Pre-compute co-prime lists for numbers from 1 to 50
33        for (int i = 1; i < 51; ++i) {
34            for (int j = 1; j < 51; ++j) {
35                if (gcd(i, j) == 1) { // i and j are co-prime
36                    coPrimeLists[i].add(j);
37                }
38            }
39        }
40      
41        this.nums = nums;
42        answers = new int[n];
43      
44        // Perform Depth-First Search from node 0 
45        dfs(0, -1, 0);
46        return answers;
47    }
48
49    private void dfs(int node, int parent, int depth) {
50        int closestCoprimeNode = -1, maxDepth = -1;
51        for (int val : coPrimeLists[nums[node]]) {
52            // Check stacks for each co-prime number to find closest ancestor
53            var stack = ancestorStacks[val];
54            if (!stack.isEmpty() && stack.peek()[1] > maxDepth) {
55                closestCoprimeNode = stack.peek()[0];
56                maxDepth = stack.peek()[1];
57            }
58        }
59        answers[node] = closestCoprimeNode;
60      
61        // Visit children
62        for (int child : graph[node]) {
63            if (child != parent) {
64                ancestorStacks[nums[node]].push(new int[] {node, depth});
65                dfs(child, node, depth + 1);
66                ancestorStacks[nums[node]].pop();
67            }
68        }
69    }
70
71    // Helper function to calculate the Greatest Common Divisor (GCD)
72    private int gcd(int a, int b) {
73        return (b == 0) ? a : gcd(b, a % b);
74    }
75}
76

C++ Solution

1#include <vector>
2#include <stack>
3#include <functional>
4#include <algorithm>
5#include <utility>
6
7class Solution {
8public:
9    std::vector<int> getCoprimes(std::vector<int>& nums, std::vector<std::vector<int>>& edges) {
10        int num_nodes = nums.size(); // Number of nodes in the graph.
11        std::vector<std::vector<int>> graph(num_nodes); // Adjacency list for the graph.
12        std::vector<std::vector<int>> coprime_numbers(51); // List to store coprimes for each number.
13        std::vector<std::stack<std::pair<int, int>>> stacks(51); // Stacks to keep track of ancestors for each number.
14
15        // Build the graph from the given edges.
16        for (auto& edge : edges) {
17            int u = edge[0], v = edge[1];
18            graph[u].push_back(v);
19            graph[v].push_back(u);
20        }
21
22        // Precompute the coprime numbers for each number from 1 to 50.
23        for (int i = 1; i < 51; ++i) {
24            for (int j = 1; j < 51; ++j) {
25                if (std::__gcd(i, j) == 1) {
26                    coprime_numbers[i].push_back(j);
27                }
28            }
29        }
30
31        std::vector<int> result(num_nodes, -1); // Initialize the result vector with -1.
32      
33        // Depth-first search function to traverse the tree and compute coprimes.
34        std::function<void(int, int, int)> depthFirstSearch = [&](int node, int parent, int depth) {
35            int closest_coprime_node = -1; // Closest coprime ancestor node.
36            int max_depth = -1; // Max depth of the coprime ancestor node.
37
38            // Find the closest coprime ancestor node.
39            for (int value : coprime_numbers[nums[node]]) {
40                auto& stack = stacks[value];
41                if (!stack.empty() && stack.top().second > max_depth) {
42                    closest_coprime_node = stack.top().first;
43                    max_depth = stack.top().second;
44                }
45            }
46
47            result[node] = closest_coprime_node; // Set closest coprime ancestor in the result.
48          
49            // Traverse the children.
50            for (int neighbor : graph[node]) {
51                if (neighbor != parent) {
52                    // Push the current node into the stack for its value before traversing its subtree.
53                    stacks[nums[node]].push({node, depth});
54                    depthFirstSearch(neighbor, node, depth + 1);
55                    // Pop the node from the stack after traversing its subtree.
56                    stacks[nums[node]].pop();
57                }
58            }
59        };
60      
61        depthFirstSearch(0, -1, 0); // Start DFS from node 0, with no parent and at depth 0.
62        return result; // Return the result vector with the closest coprime ancestors.
63    }
64};
65

Typescript Solution

1// Importing Array utility from lodash for gcd function
2import _ from 'lodash';
3
4// Function to build a graph from given edges
5function buildGraph(edges: number[][]): number[][] {
6    const graph: number[][] = [];
7    for (const edge of edges) {
8        const u = edge[0], v = edge[1];
9        if (!graph[u]) graph[u] = [];
10        if (!graph[v]) graph[v] = [];
11        graph[u].push(v);
12        graph[v].push(u);
13    }
14    return graph;
15}
16
17// Function to compute and store coprime numbers for each integer from 1 to 50
18function computeCoprimeNumbers(): number[][] {
19    const coprimeNumbers: number[][] = [];
20    for (let i = 1; i <= 50; i++) {
21        coprimeNumbers[i] = [];
22        for (let j = 1; j <= 50; j++) {
23            if (_.gcd(i, j) === 1) {
24                coprimeNumbers[i].push(j);
25            }
26        }
27    }
28    return coprimeNumbers;
29}
30
31// The main function which computes the closest coprime ancestor for each node
32function getCoprimes(nums: number[], edges: number[][]): number[] {
33    const numNodes: number = nums.length;
34    const graph: number[][] = buildGraph(edges);
35    const coprimeNumbers: number[][] = computeCoprimeNumbers();
36    const stacks: Array<StackItem>[] = Array.from({ length: 51 }, () => []);
37
38    const result: number[] = nums.map(() => -1);
39
40    // Defines a type for stack items which includes node index and depth
41    type StackItem = { node: number; depth: number; };
42
43    // Recursive Depth-First Search (DFS) function
44    function depthFirstSearch(node: number, parent: number, depth: number): void {
45        let closestCoprimeNode: number = -1;
46        let maxDepth: number = -1;
47
48        // Find the closest coprime ancestor for the current node
49        for (const value of coprimeNumbers[nums[node]]) {
50            const stack = stacks[value];
51            if (stack.length && stack[stack.length - 1].depth > maxDepth) {
52                closestCoprimeNode = stack[stack.length - 1].node;
53                maxDepth = stack[stack.length - 1].depth;
54            }
55        }
56
57        // Save the closest coprime ancestor for the current node
58        result[node] = closestCoprimeNode;
59
60        // Push the current node into the stack representing its numeric value before DFS on children
61        stacks[nums[node]].push({ node, depth });
62
63        // DFS on children nodes
64        for (const neighbor of graph[node]) {
65            if (neighbor !== parent) {
66                depthFirstSearch(neighbor, node, depth + 1);
67            }
68        }
69
70        // Pop the current node after DFS on its subtree is complete
71        stacks[nums[node]].pop();
72    }
73
74    depthFirstSearch(0, -1, 0);
75
76    return result;
77}
78
79// eslint-disable-next-line @typescript-eslint/no-export-default
80export default getCoprimes;
81

Time and Space Complexity

Time Complexity

The time complexity of the given code mainly consists of three parts: the construction of graph g, the creation of co-prime list f, and the depth-first search function dfs.

  1. Graph Construction: It involves iterating through the edges list once to build the graph g with adjacency lists, which takes O(E), where E is the number of edges.

  2. Co-Prime List Construction: The nested loop to initialize the f dictionary runs through all pairs of numbers from 1 to 50, which takes O(50^2), effectively O(1) since it's a constant.

  3. Depth-First Search (DFS):

    • The DFS function dfs is called once for each node in the nums list (tree). Since it's a tree, there are N-1 edges for N nodes, meaning the DFS will handle each node exactly once. Therefore, it has a linear relationship with the number of nodes, giving O(N).
    • Within the DFS, we search for the most recent coprime ancestor. In the worst case, we might have to iterate through all ancestors stored in the stacks for each node. In the worst-case scenario, we could have O(N) ancestors for each node, due to the fact that each node can have a different value in nums leading to distinct keys in stks.
    • The append and pop operations on stks are O(1) each, but as they may occur for every node, their inclusion in the DFS traversal does not change the overall time complexity.

Combining these parts, we get O(E) + O(1) + O(N) * O(N), but since we know that E is effectively O(N), we can simplify this to O(N^2) for the worst-case time complexity due to the DFS section being the dominating factor.

Space Complexity

The space complexity is determined by the graph g, the co-prime list f, the stacks stks, and the DFS call stack.

  1. Graph g: Takes O(E) space, which is O(N) for a tree.

  2. Co-Prime List f: Takes O(50 * 50) space, which simplifies to O(1) since it's constant.

  3. Stacks stks: In the worst case, the stacks can grow to store an ancestor for each value in nums, meaning it could take O(N * 50) space. However, considering each value in nums may store a list of ancestors up to O(N), it could require up to O(N^2) in the worst case.

  4. DFS Call Stack: The recursive calls to DFS would use O(N) space since a tree can have at most N calls on the call stack at any given time in the case of a skewed tree.

Combining the above, we get a total worst-case space complexity of O(N^2) due to the stks being the dominating factor.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫