洛谷4427 [BJOI2018]求和

题意

一颗 $N$ 点的树，有 $M$ 组查询 $(i,j,k)$ ，询问 $i\rightarrow j$ 所有节点深度的 $k$ 次方和。

$N,M\le 300000 \ , \ K\le 50$ 。

题解

因为 $K$ 很小，所以每个点的所有权值都可以预处理出来。

本来想写树剖的，但其实树上差分就够了。

$$ans=v[i]+v[j]-v[lca(i,j)]-v[fa[lca(i,j)]]$$

#include<bits/stdc++.h>
#define ll long long
#define ha 998244353

using namespace std;

inline ll read()
{
    char ch=getchar(); ll f=1,x=0;
    while (ch<'0' || ch>'9') { if (ch=='-') f=-1; ch=getchar(); }
    while (ch>='0' && ch<='9') { x=x*10+ch-'0'; ch=getchar(); }
    return f*x;
}

struct Edge {
    ll next,to;
} edge[600005];
const ll N=300005;
ll cnt,head[N],n,m,a,b,c,fa[N][20],deep[N],lg[N],v[N][55];

inline void add(ll u,ll v)
{
    edge[++cnt].to=v;
    edge[cnt].next=head[u];
    head[u]=cnt;
}

void dfs(ll x,ll f,ll dep)
{
    deep[x]=dep;
    v[x][1]=dep;
    for (ll i=2;i<=50;i++) v[x][i]=(v[x][i-1]*dep)%ha;
    for (ll i=1;i<=50;i++) v[x][i]=(v[x][i]+v[f][i])%ha;
    for (ll i=1;i<=lg[deep[x]];i++) fa[x][i]=fa[fa[x][i-1]][i-1];
    for (ll i=head[x];i;i=edge[i].next)
    {
        ll y=edge[i].to;
        if (y==f) continue;
        fa[y][0]=x;
        dfs(y,x,dep+1);
    }
}

inline ll lca(ll u,ll v)
{
    if (deep[u]<deep[v]) u^=v^=u^=v;
    while (deep[u]>deep[v]) u=fa[u][lg[deep[u]-deep[v]]];
    if (u==v) return u;
    for (ll i=lg[deep[u]];~i;i--) if (fa[u][i]^fa[v][i]) u=fa[u][i],v=fa[v][i];
    return fa[u][0];
}

signed main()
{
    n=read();
    for (ll i=1;i<n;i++)
    {
        a=read(); b=read();
        add(a,b);
        add(b,a);
    }
    for (ll i=1;i<=n;i++)
    {
        lg[i]=lg[i-1];
        if (1<<(lg[i]+1)==i) lg[i]++;
    }
    dfs(1,0,0);
    m=read();
    for (ll i=1;i<=m;i++)
    {
        a=read(); b=read(); c=read();
        ll LCA=lca(a,b);
        printf("%lld\n",(v[a][c]+v[b][c]-v[LCA][c]-v[fa[LCA][0]][c]+2*ha)%ha);
    }
    return 0;
}