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 nwhere 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

Efficient Solutions by OpHaxor:

    public int[] getCoprimes(int[] nums, int[][] edges) {
        int n = nums.length;
        List<List<Integer>> adj = new ArrayList<>(n);
        for (int i = 0; i < n; i++) adj.add(new ArrayList<>());
        for (int[] e : edges) {
            adj.get(e[0]).add(e[1]);
            adj.get(e[1]).add(e[0]);
        }
        ans = new int[n];
        for (int i = 1; i <= 50; i++) {
            for (int j = 1; j <= 50; j++) {
                coprime[i][j] = gcd(i, j) == 1;
            }
        }
        stacks = new ArrayList<>(51);
        for (int i = 0; i <= 50; i++) stacks.add(new ArrayDeque<int[]>());
        dfs(nums, adj, 0, -1, 0);
        return ans;
    }
    
    int[] ans;
    boolean[][] coprime = new boolean[51][51];
    // id, depth
    List<ArrayDeque<int[]>> stacks;
    
    void dfs(int[] nums, List<List<Integer>> adj, int u, int p, int d) {
        int me = nums[u];
        int[] best = new int[]{-1, -1};
        for (int other = 1; other <= 50; other++) {
            if (!coprime[me][other]) continue;
            if (!stacks.get(other).isEmpty()) {
                int[] anc = stacks.get(other).peekLast();
                if (anc[1] > best[1]) {
                    best[0] = anc[0];
                    best[1] = anc[1];
                }
            }
        }
        ans[u] = best[0];
        stacks.get(me).addLast(new int[]{u, d});
        for (int v : adj.get(u)) {
            if (v == p) continue;
            dfs(nums, adj, v, u, d+1);
        }
        stacks.get(me).removeLast();
    }
    
    int gcd(int a, int b) {
        return a == 0 ? b : gcd(b % a, a);
    }
def getCoprimes(self, nums: List[int], edges: List[List[int]]) -> List[int]:
        length = len(nums)

        graph = [[] for i in range(length)]
        for first, second in edges:
            graph[first].append(second)
            graph[second].append(first)

        answers = [-1 for _ in range(length)]
        path = defaultdict(list)

        def dfs(vertex: int = 0, parent: int = -1, depth: int = 0):
            best_depth = -1
            for value, parents in path.items():
                if math.gcd(value, nums[vertex]) == 1 and parents:
                    if parents[-1][1] > best_depth:
                        answers[vertex] = parents[-1][0]
                        best_depth = parents[-1][1]

            path[nums[vertex]].append((vertex, depth))

            for adjacent in graph[vertex]:
                if adjacent != parent:
                    dfs(adjacent, vertex, depth + 1)

            path[nums[vertex]].pop()

        dfs()
        
        return answers
vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {
        int n = nums.size();
        vector<vector<int>> neighbours(n);
        for(const auto& edge : edges){
            neighbours[edge[0]].push_back(edge[1]);
            neighbours[edge[1]].push_back(edge[0]);
        }
        vector<int> level(n, 0);
        // level order traversal
        queue<int> q;
        q.push(0);
        vector<bool> visited(n, false);
        visited[0] = true;
        int depth = 0;
        while(!q.empty()){
            int size = q.size();
            while(size--){
                int cur = q.front();
                q.pop();
                level[cur] = depth;
                for(int next : neighbours[cur]){
                    if(!visited[next]){
                        visited[next] = true;
                        q.push(next);
                    }
                }
            }
            depth++;
        }
        
        vector<int> ans(n, -1);
        visited = vector<bool>(n, false);
        vector<stack<int>> lastSeen(51);
        dfs(neighbours, nums, level, 0, visited, lastSeen, ans);
        return ans;
    }
private:
    int gcd(int a, int b){
        if(a == 0){
            return b;
        }
        if(b == 0){
            return a;
        }
        return gcd(b, a%b);
    }
    
    void dfs(const vector<vector<int>>& neighbours, const vector<int>& nums, const vector<int>& level, int cur, vector<bool>& visited, vector<stack<int>>& lastSeen, vector<int>& ans){
        for(int i = 1; i < lastSeen.size(); i++){
            if(!lastSeen[i].empty()){
                int j = lastSeen[i].top();
                if(gcd(nums[cur], nums[j]) == 1){
                    if(ans[cur] == -1 || level[ans[cur]] < level[j]){
                        ans[cur] = j;
                    }
                }
            }
        }
        lastSeen[nums[cur]].push(cur);
        visited[cur] = true;
        for(int next : neighbours[cur]){
            if(!visited[next]){
                dfs(neighbours, nums, level, next, visited, lastSeen, ans);
            }
        }
        lastSeen[nums[cur]].pop();
    }

Leave a reply