LeetCode 1766. Tree of Coprimes

題目

題意： 對於樹上的每個節點，找出與它的值互質的最近的祖先節點。

題解： 由於節點上的值在1-50之間，所以算互質很好算，事先算法。然後就是深度優先遍歷樹的時候維護路徑上的節點的位置，利用1-50這個小範圍，快速找到與當前節點值互質的值出現在哪些位置上

struct Node
{
    int value;
    int next;
}edge[500005];

int head[100005];
int pos;
void init()
{
    memset(head,-1,sizeof(head));
    pos=0;
}

void add(int x,int y)
{
    edge[pos].value = y;
    edge[pos].next = head[x];
    head[x] = pos++;
}

int gcd(int x, int y)
{
    if(x<y)
        swap(x,y);
    
    return x%y==0 ? y : gcd(x%y, y); 
}

class Solution {
public:
    vector<int> p[55];
    int s[100005];
    int s_p;
    vector<int> a[55];
    int vis[100005];
    vector<int> ans;
    vector<int> getCoprimes(vector<int>& nums, vector<vector<int>>& edges) {
        
        init();
        for(int i=0;i<edges.size();i++)
        {
            add(edges[i][0], edges[i][1]);
            add(edges[i][1], edges[i][0]);
        }
        
        for(int i=1;i<=50;i++)
        {
            for(int j=1;j<=50;j++)
            {
                if(gcd(i,j)==1)
                {
                    p[i].push_back(j);
                }
            }
        }
        
        memset(vis,0,sizeof(vis));
        s_p=0;
        
        for(int i=0;i<nums.size();i++)
        {
            ans.push_back(0);
        }
        s[s_p] = 0;
        a[nums[0]].push_back(0);
        vis[0]=1;
        ans[0]=-1;
        dfs(0, nums);
        
        return ans;
        
        
    }
    
    void dfs(int root, vector<int>& nums)
    {
        
        
        for(int i=head[root]; i!=-1; i=edge[i].next)
        {
            int v = edge[i].value;
            if(vis[v]==1)
                continue;
            
            int res = -1;
            for(int j=0;j<p[nums[v]].size();j++)
            {
                if(a[p[nums[v]][j]].size()>0)
                {
                    res = max(res, a[p[nums[v]][j]][a[p[nums[v]][j]].size()-1]);
                }
            }
            
            if(res!=-1)
                ans[v] = s[res];
            else
                ans[v] = -1;
            
            vis[v]=1;
            
            s[++s_p] = v;
            a[nums[v]].push_back(s_p);
            
            dfs(v, nums);
            
            s_p--;
            a[nums[v]].pop_back();
        }
    }
};