CodeForces - 1101D GCD Counting [트리 지름]

D. GCD Counting
time limit per test
4.5 seconds
memory limit per test
256 megabytes
input
standard input
output
standard output
You are given a tree consisting of nn vertices. A number is written on each vertex; the number on vertex ii is equal to aiai.
Let's denote the function g(x,y)g(x,y) as the greatest common divisor of the numbers written on the vertices belonging to the simple path from vertex xx to vertex yy (including these two vertices). Also let's denote dist(x,y)dist(x,y) as the number of vertices on the simple path between vertices xx and yy, including the endpoints. dist(x,x)=1dist(x,x)=1 for every vertex xx.
Your task is calculate the maximum value of dist(x,y)dist(x,y) among such pairs of vertices that g(x,y)>1g(x,y)>1.
Input
The first line contains one integer nn — the number of vertices (1≤n≤2⋅105)(1≤n≤2⋅105).
The second line contains nn integers a1a1, a2a2, ..., anan (1≤ai≤2⋅105)(1≤ai≤2⋅105) — the numbers written on vertices.
Then n−1n−1 lines follow, each containing two integers xx and yy (1≤x,y≤n,x≠y)(1≤x,y≤n,x≠y) denoting an edge connecting vertex xx with vertex yy. It is guaranteed that these edges form a tree.
Output
If there is no pair of vertices x,yx,y such that g(x,y)>1g(x,y)>1, print 00. Otherwise print the maximum value of dist(x,y)dist(x,y) among such pairs.
 
 
한 사슬의 모든 수의 gcd>1은 그들이 같은 질인자를 가지고 있다는 것을 이해할 수 있기 때문에 모든 질수를 일일이 열거하고 나무의 직경을 뛰기만 하면 된다.
1부터 2e5까지의 모든 질인자를 선형으로 선별한 다음에 각 질수에 대응하는 점을 찾아내고 매거한 다음에 dfs를 뛴다.

#include "bits/stdc++.h"
using namespace std;
int a[200004];
vectorisp;//  
setprim[200004];//       
vectorxxx[200004];//        
bool vis[200004];
bool vs[200004];
struct edge
{
    int v,nxt;
}g[400004];
int head[400004];
int cnt;
int n;
void addedge(int u,int v)
{
    g[cnt].v=v;
    g[cnt].nxt=head[u];
    head[u]=cnt;
    ++cnt;
}
int ans=0;
int tu;
void dfs(int u,int pre,int gc,int num,int ok)
{
    if(!ok&&vs[u])return ;
    vs[u]=1;
    for (int i = head[u]; i !=-1 ; i=g[i].nxt) {
        int v=g[i].v;
        if(v!=pre&&(prim[a[v]].find(gc)!=prim[a[v]].end()))
            dfs(v,u,gc,num+1,ok);
    }
    if(num>ans)
    {
        ans=num;
        tu=u;
    }
}
int main()
{
    cin>>n;
    cnt=0;
    memset(head,-1, sizeof(head));
    memset(vis,0, sizeof(vis));
    memset(vs,0, sizeof(vs));
    for (int i = 2; i <= 200000 ; ++i) {// 
        if(!vis[i]) {
            isp.push_back(i);
            for (int j = i ; j <= 200000; j += i) {
                if(j!=i)vis[j] = 1;
                prim[j].insert(i);
            }
        }
    }
    int ok=0;
    int maxn=0;
    for (int i = 1; i <= n; ++i) {
        scanf("%d",&a[i]);
        for (auto j = prim[a[i]].begin(); j != prim[a[i]].end(); ++j) {//        
            xxx[*j].push_back(i);
            maxn=max(maxn,*j);
        }
        if(a[i]!=1)ok=1;

    }
    for (int i = 0; i < n-1; ++i) {
        int x,y;
        scanf("%d%d",&x,&y);
        addedge(x,y);
        addedge(y,x);
    }

    int ANS=0;
    for (int i = 2; i <= maxn; ++i) {//    
        if(!vis[i])
        {
            memset(vs,0, sizeof(vs));
            for (int j = 0; j < xxx[i].size(); ++j) {//      ，                  
                ans=0;
                int xxxx=tu;
                dfs(xxx[i][j],-1,i,1,0);
                if(xxxx!=tu){
                    dfs(tu,-1,i,1,1);
                    ANS=max(ANS,ans);
                }

            }
        }
    }
    if(ok)
    printf("%d
",ANS);
    else puts("0");
}