Cod sursa(job #2959621)

Utilizator MihaiBunBunget Mihai MihaiBun Data 1 ianuarie 2023 19:47:19
Problema Lowest Common Ancestor Scor 0
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.38 kb
#include <fstream>
#include <vector>
#define N 100003
using namespace std;
ifstream f("sumtree.in");
ofstream gg("sumtree.out");
int n,i,u,v,j,ha,nr;
long long sol;
int val[N],viz[N],e[2*N],first[N],niv[N],nrarb[N],sump[N],sumval[N],last[N];
int lg[21],rmq[21][2*N],ta[N];
vector<int> g[N];

void tree(int nod, int nivel)
{
    viz[nod] = 1;
    nr++;
    e[nr] = nod;
    first[nod] = nr;
    niv[nod] = nivel;
    nrarb[nod] = 1;
    sump[nod] = 0;
    sumval[nod] = sumval[ta[nod]]+val[nod];
    int h;
    for(h = 0; h < g[nod].size(); h++)
        if(viz[g[nod][h]]==0)
    {
        ta[g[nod][h]] = nod;
        tree(g[nod][h],nivel+1);
        nrarb[nod] += nrarb[g[nod][h]];
        sump[nod] += nrarb[g[nod][h]]*nrarb[g[nod][h]];
        nr++;
        e[nr] = nod;
    }

}

void makermq()
{
    int l,i,j;
    lg[1]=0;
	for (i=2;i<=nr;i++)
		lg[i]=lg[i/2]+1;

	for (i=1;i<=nr;i++)
		rmq[0][i]=i;

	for (i=1; (1 << i) <=nr;i++)
		for (j=1;j <= nr - (1 << i)+1;j++)
		{
		l=1<<(i-1);
		if(niv[e[rmq[i-1][j]]] < niv[e[rmq[i-1][j+l]]])
            rmq[i][j] = rmq[i-1][j];
        else rmq[i][j] = rmq[i-1][j+l];
		}
}

int lca(int x, int y)
{
    int diff,l,sh;
    if(x > y) swap(x,y);
    diff=y-x+1;
    l=lg[diff];
    sh=diff - (1<<l);
    if(niv[e[rmq[l][x]]] < niv[e[rmq[l][x+sh]]]) return e[rmq[l][x]];
    else return e[rmq[l][x+sh]];
}

int gcd(int x, int y)
{
    int r = x%y;
    while(r!=0)
    {
        x = y;
        y = r;
        r = x%y;
    }
    return y;
}

int main()
{
    f >> n;
    for(i = 1; i <= n; i++)
        f >> val[i];
    for(i = 1; i < n; i++)
    {
        f >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    nr = 0;
    ta[1] = 0;
    tree(1,0);
    if(n <= 1000)
    {
        makermq();
        sol = 0;
        for(i = 1; i < n; i++)
            for(j = i+1; j <= n; j++)
            if(gcd(val[i],val[j])>1)
        {
            ha = lca(first[i],first[j]);
            sol += sumval[i]+sumval[j]-2*sumval[ha]+val[ha];
        }
        gg << sol;
    }
    else
    {
        if(val[1]) sol = 0;
        else
        {
            sol = 0;
            for(i = 1; i <= n; i++)
                sol += 1LL*(n-1+1LL*(nrarb[i]-1)*(n-nrarb[i])+(1LL*(nrarb[i]-1)*(nrarb[i]-1)-sump[i])/2)*val[i];
        }
        gg << sol;
    }
    return 0;
}