1766

1766. Tree of Coprimes (Hard)

There is a tree (i.e., a connected, undirected graph that has no cycles) consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges. Each node has a value associated with it, and the root of the tree is node 0.

To represent this tree, you are given an integer array nums and a 2D array edges. Each nums[i] represents the ith node's value, and each edges[j] = [uj, vj] represents an edge between nodes uj and vj in the tree.

Two values x and y are coprime if gcd(x, y) == 1 where gcd(x, y) is the greatest common divisor of x and y.

An ancestor of a node i is any other node on the shortest path from node i to the root. A node is not considered an ancestor of itself.

Return an array ans of size n, where ans[i] is the closest ancestor to node i such that nums[i] and nums[ans[i]] are coprime, or -1 if there is no such ancestor.

 

Example 1:

Input: nums = [2,3,3,2], edges = [[0,1],[1,2],[1,3]]
Output: [-1,0,0,1]
Explanation: In the above figure, each node's value is in parentheses.
- Node 0 has no coprime ancestors.
- Node 1 has only one ancestor, node 0. Their values are coprime (gcd(2,3) == 1).
- Node 2 has two ancestors, nodes 1 and 0. Node 1's value is not coprime (gcd(3,3) == 3), but node 0's
  value is (gcd(2,3) == 1), so node 0 is the closest valid ancestor.
- Node 3 has two ancestors, nodes 1 and 0. It is coprime with node 1 (gcd(3,2) == 1), so node 1 is its
  closest valid ancestor.

Example 2:

Input: nums = [5,6,10,2,3,6,15], edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]
Output: [-1,0,-1,0,0,0,-1]

 

Constraints:

• nums.length == n
• 1 <= nums[i] <= 50
• 1 <= n <= 105
• edges.length == n - 1
• edges[j].length == 2
• 0 <= uj, vj < n
• uj != vj

Related Topics:
Math, Tree, Depth-first Search, Breadth-first Search

Solution 1. DFS

Complexity Analysis

Collecting all the unique numbers takes O(N) time and O(50) = O(1) space.

Generating the cops takes O(50 * 50) = O(1) time.

Generating the graph G takes O(N) time and O(N) space.

For each DFS invocation, traversing all the cops to get the maximum level takes O(50) = O(1) time.

The maximum depth of the DFS is O(N), so this take O(N) time and O(N) space.

Overall, this solution takes O(N) time and O(N) space.

// OJ: https://leetcode.com/problems/tree-of-coprimes/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
// Ref: https://leetcode.com/problems/tree-of-coprimes/discuss/1074770/C%2B%2B-Save-the-cops
class Solution {
    pair<int, int> maxLevel[51]; // level, index
    vector<int> cops[51], ans;
    vector<vector<int>> G;
    void dfs(vector<int> &A, int i, int parent = 0, int lv = 0) {
        int maxLv = -1;
        for (int cop : cops[A[i]]) { // traverse each co-prime of `A[i]` and find the one with the maximum level
            auto [level, index] = maxLevel[cop];
            if (level > maxLv) {
                maxLv = level;
                ans[i] = index;
            }
        }
        auto prev = maxLevel[A[i]];
        maxLevel[A[i]] = { lv, i }; // update the current level as the maximum level of A[i]
        for (int j : G[i]) {
            if (j != parent) dfs(A, j, i, lv + 1);
        }
        maxLevel[A[i]] = prev; // backtrack
    }
public:
    vector<int> getCoprimes(vector<int>& A, vector<vector<int>>& E) {
        unordered_set<int> s(begin(A), end(A));  // collect all the unique numbers
        for (int n : s) {
            for (int m : s) {
                if (gcd(n, m) == 1) cops[n].push_back(m); // for each number, save all of its co-primes into `cops`
            }
        }
        G.assign(A.size(), {});
        for (auto &e : E) {
            int u = e[0], v = e[1];
            G[u].push_back(v);
            G[v].push_back(u);
        }
        ans.assign(A.size(), -1);
        memset(maxLevel, -1, sizeof(maxLevel));
        dfs(A, 0);
        return ans;
    }
};