upc 6744 求和(倍增求LCA)

#include<stdio.h>
#include<iostream>
#include<string.h>
#include<algorithm>
#include<vector>
#include<math.h>
using namespace std;
#define ll long long
const ll mod=998244353;
#define maxn 300050
int f[maxn][20];
ll sum[maxn][55];
int w[maxn];
int head[maxn];
int tot,N;
struct Edge
{
    int to;
    int next;
} E[maxn*2];
void init(int n)
{
    memset(head,-1,sizeof(head));
    memset(f,0,sizeof(f));
    memset(sum,0,sizeof(sum));
    memset(w,0,sizeof(w));
    tot=0;
    N=floor(log(n+0.0)/log(2));
    return;
}
void addedge(int u,int v)
{
    E[tot].to=v;
    E[tot].next=head[u];
    head[u]=tot++;
    E[tot].to=u;
    E[tot].next=head[v];
    head[v]=tot++;
    return;
}
ll qpow(int u,int k)
{
    ll num=(ll)u;
    ll ans=1;
    while(k)
    {
        if(k&1)
            ans=ans*num%mod;
        k>>=1;
        num=num*num%mod;
    }
    return ans;
}
void dfs(int su,int u,int e)
{
    w[u]=e;
    for(int i=1; i<=55; i++)
    {
        sum[u][i]=sum[su][i]+qpow(w[u],i)%mod;
    }
    for(int i=1; i<=N; i++)
    {
        f[u][i]=f[f[u][i-1]][i-1];
    }
    for(int i=head[u]; i!=-1; i=E[i].next)
    {
        if(E[i].to==su)continue;
        f[E[i].to][0]=u;
        dfs(u,E[i].to,e+1);
    }
    return;
}
int lca(int a,int b)
{
    if(w[a]>w[b])swap(a,b);
    for(int i=0; i<=N; i++)
    {
        if((1<<i)&w[b]-w[a])
        {
            b=f[b][i];
        }
    }
    if(b==a)return a;
    for(int i=N; i>=0; i--)
    {
        if(f[a][i]!=f[b][i])
        {
            a=f[a][i];
            b=f[b][i];
        }
    }
    return f[a][0];

}
int main()
{
    int n,m,a,b,c;
    scanf("%d",&n);
    init(n);
    for(int i=1; i<n; i++)
    {
        scanf("%d%d",&a,&b);
        addedge(a,b);
    }
    dfs(1,1,0);
    scanf("%d",&m);
    int fa;
    for(int i=1; i<=m; i++)
    {
        scanf("%d%d%d",&a,&b,&c);
        fa=lca(a,b);
        printf("%lld\n",((sum[a][c]+sum[b][c]-sum[fa][c]-sum[f[fa][0]][c])%mod+mod)%mod);
    }
}

猜你喜欢

转载自blog.csdn.net/zyy_1998/article/details/81587909
UPC
今日推荐