1766. Tree of Coprimes

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.

class Solution {
    vector<vector<int>> g;
    vector<int> res;
    unordered_map<int,vector<int>> conn;
    void helper(int node, int p, unordered_map<int,vector<pair<int,int>>>& his, vector<int>& nums, int lvl) {
        int malvl = -1;
        for(auto con : conn[nums[node]]) {
            if(!his[con].empty() and his[con].back().second > malvl) {
                malvl = his[con].back().second;
                res[node] = his[con].back().first;
            }
        }
        
        his[nums[node]].push_back({node,lvl});
        
        for(auto near : g[node]) {
            if(near != p)
                helper(near, node, his, nums, lvl+1);
        }
        
        his[nums[node]].pop_back();
        
    }
public:
    vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {
        int n = nums.size();
        g = vector<vector<int>>(n);
        res = vector<int>(n,-1);
        unordered_set<int> s(nums.begin(), nums.end());
        for(auto n1 : s)
            for(auto n2 : s)
                if(__gcd(n1,n2)==1)
                    conn[n1].push_back(n2);
        
        for(auto& edge : edges) {
            g[edge[0]].push_back(edge[1]);
            g[edge[1]].push_back(edge[0]);
        }
        
        unordered_map<int,vector<pair<int,int>>> tmp;
        helper(0,-1,tmp,nums,0);
        
        return res;
    }
};