【刷题】BZOJ 5293 [Bjoi2018]求和

Description

master 对树上的求和非常感兴趣。他生成了一棵有根树，并且希望多次询问这棵树上一段路径上所有节点深度的k

次方和，而且每次的k 可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给

了pupil，但pupil 并不会这么复杂的操作，你能帮他解决吗？

Input

第一行包含一个正整数n ，表示树的节点数。

之后n-1 行每行两个空格隔开的正整数i,j ，表示树上的一条连接点i 和点j 的边。

之后一行一个正整数m ，表示询问的数量。

之后每行三个空格隔开的正整数i,j,k ，表示询问从点i 到点j 的路径上所有节点深度的k 次方和。

由于这个结果可能非常大，输出其对998244353 取模的结果。

树的节点从1 开始标号，其中1 号节点为树的根。

Output

对于每组数据输出一行一个正整数表示取模后的结果。

1≤n,m≤300000,1≤k≤50

Sample Input

5

1 2

1 3

2 4

2 5

2

1 4 5

5 4 45

Sample Output

33

503245989

说明

样例解释

以下用d(i) 表示第i 个节点的深度。

对于样例中的树，有d(1)=0,d(2)=1,d(3)=1,d(4)=2,d(5)=2。

因此第一个询问答案为(2^5 + 1^5 + 0^5) mod 998244353 = 33

第二个询问答案为(2^45 + 1^45 + 2^45) mod 998244353 = 503245989。

Solution

这题就一水题，指数小于等于50，直接对于每一个指数都维护树上两点权值和

开始写了一个树剖，交洛谷加O2过了，然后发现跑得有点慢，就又写了一个差分

这是差分的代码：

#include<bits/stdc++.h>
#define ui unsigned int
#define ll long long
#define db double
#define ld long double
#define ull unsigned long long
const int MAXN=300000+10,Mod=998244353;
int n,q,e,to[MAXN<<1],nex[MAXN<<1],beg[MAXN],Jie[MAXN][20],dep[MAXN];
ll Sum[MAXN][51];
template<typename T> inline void read(T &x)
{
	T data=0,w=1;
	char ch=0;
	while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();
	if(ch=='-')w=-1,ch=getchar();
	while(ch>='0'&&ch<='9')data=((T)data<<3)+((T)data<<1)+(ch^'0'),ch=getchar();
	x=data*w;
}
template<typename T> inline void write(T x,char ch='\0')
{
	if(x<0)putchar('-'),x=-x;
	if(x>9)write(x/10);
	putchar(x%10+'0');
	if(ch!='\0')putchar(ch);
}
template<typename T> inline void chkmin(T &x,T y){x=(y<x?y:x);}
template<typename T> inline void chkmax(T &x,T y){x=(y>x?y:x);}
template<typename T> inline T min(T x,T y){return x<y?x:y;}
template<typename T> inline T max(T x,T y){return x>y?x:y;}
inline void insert(int x,int y)
{
	to[++e]=y;
	nex[e]=beg[x];
	beg[x]=e;
}
inline ll qexp(ll a,ll b)
{
	ll res=1;
	while(b)
	{
		if(b&1)res=res*a%Mod;
		a=a*a%Mod;
		b>>=1;
	}
	return res;
}
inline void dfs1(int x,int f)
{
	Jie[x][0]=f;dep[x]=dep[f]+1;
	for(register int i=beg[x];i;i=nex[i])
		if(to[i]==f)continue;
		else dfs1(to[i],x);
	for(register int i=1;i<=50;++i)Sum[x][i]=qexp(dep[x],i);
}
inline void dfs2(int x)
{
	for(register int i=1;i<=50;++i)(Sum[x][i]+=Sum[Jie[x][0]][i])%=Mod;
	for(register int i=beg[x];i;i=nex[i])
		if(to[i]==Jie[x][0])continue;
		else dfs2(to[i]);
}
inline void init()
{
	dfs1(1,0);
	dfs2(1);
	for(register int j=1;j<19;++j)
		for(register int i=1;i<=n;++i)Jie[i][j]=Jie[Jie[i][j-1]][j-1];
}
inline int LCA(int u,int v)
{
	if(dep[u]<dep[v])std::swap(u,v);
	if(dep[u]>dep[v])
		for(register int i=19;i>=0;--i)
			if(dep[Jie[u][i]]>=dep[v])u=Jie[u][i];
	if(u==v)return u;
	for(register int i=19;i>=0;--i)
		if(Jie[u][i]!=Jie[v][i])u=Jie[u][i],v=Jie[v][i];
	return Jie[u][0];
}
int main()
{
	read(n);
	for(register int i=1;i<n;++i)
	{
		int u,v;read(u);read(v);
		insert(u,v);insert(v,u);
	}
	dep[0]=-1;
	init();
	read(q);
	while(q--)
	{
		int u,v,k;read(u);read(v);read(k);
		int lca=LCA(u,v);
		write((Sum[u][k]+Sum[v][k]-Sum[lca][k]+Mod-Sum[Jie[lca][0]][k]+Mod)%Mod,'\n');
	}
	return 0;
}

这是树剖的代码（开O2能过，不知道不开O2或者其他OJ上能不能过）：

#include<bits/stdc++.h>
#define ui unsigned int
#define ll long long
#define db double
#define ld long double
#define ull unsigned long long
const int MAXN=300000+10,Mod=998244353;
int n,q,e,to[MAXN<<1],nex[MAXN<<1],beg[MAXN],st[MAXN],ed[MAXN],fa[MAXN],dep[MAXN],val[MAXN],hson[MAXN],size[MAXN],cnt,top[MAXN];
#define Mid ((l+r)>>1)
#define ls rt<<1
#define rs rt<<1|1
#define lson ls,l,Mid
#define rson rs,Mid+1,r
struct Segment_Tree{
    ll Sum[MAXN<<2][51];
    inline ll qexp(ll a,ll b)
    {
        ll res=1;
        while(b)
        {
            if(b&1)res=res*a%Mod;
            a=a*a%Mod;
            b>>=1;
        }
        return res;
    }
    inline void PushUp(int rt)
    {
        for(register int i=1;i<=50;++i)Sum[rt][i]=(Sum[ls][i]+Sum[rs][i])%Mod;
    }
    inline void Build(int rt,int l,int r)
    {
        if(l==r)
            for(register int i=1;i<=50;++i)Sum[rt][i]=qexp(val[l],i);
        else
        {
            Build(lson);Build(rson);
            PushUp(rt);
        }
    }
    inline ll Query(int rt,int l,int r,int L,int R,int k)
    {
        if(L<=l&&r<=R)return Sum[rt][k];
        else
        {
            ll res=0;
            if(L<=Mid)(res+=Query(lson,L,R,k))%=Mod;
            if(R>Mid)(res+=Query(rson,L,R,k))%=Mod;
            return res;
        }
    }
};
Segment_Tree T;
#undef Mid
#undef ls
#undef rs
#undef lson
#undef rson
template<typename T> inline void read(T &x)
{
    T data=0,w=1;
    char ch=0;
    while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();
    if(ch=='-')w=-1,ch=getchar();
    while(ch>='0'&&ch<='9')data=((T)data<<3)+((T)data<<1)+(ch^'0'),ch=getchar();
    x=data*w;
}
template<typename T> inline void write(T x,char ch='\0')
{
    if(x<0)putchar('-'),x=-x;
    if(x>9)write(x/10);
    putchar(x%10+'0');
    if(ch!='\0')putchar(ch);
}
template<typename T> inline void chkmin(T &x,T y){x=(y<x?y:x);}
template<typename T> inline void chkmax(T &x,T y){x=(y>x?y:x);}
template<typename T> inline T min(T x,T y){return x<y?x:y;}
template<typename T> inline T max(T x,T y){return x>y?x:y;}
inline void insert(int x,int y)
{
    to[++e]=y;
    nex[e]=beg[x];
    beg[x]=e;
}
inline void dfs1(int x,int f)
{
    int res=0;
    dep[x]=dep[f]+1;size[x]=1;fa[x]=f;
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==f)continue;
        else
        {
            dfs1(to[i],x);
            size[x]+=size[to[i]];
            if(size[to[i]]>res)res=size[to[i]],hson[x]=to[i];
        }
}
inline void dfs2(int x,int tp)
{
    top[x]=tp;st[x]=++cnt;val[cnt]=dep[x];
    if(hson[x])dfs2(hson[x],tp);
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==fa[x]||to[i]==hson[x])continue;
        else dfs2(to[i],to[i]);
    ed[x]=cnt;
}
inline void init()
{
    dfs1(1,0);
    dfs2(1,1);
    T.Build(1,1,n);
}
inline ll Getans(int u,int v,int k)
{
    ll res=0;
    while(top[u]!=top[v])
    {
        if(dep[top[u]]<dep[top[v]])std::swap(u,v);
        (res+=T.Query(1,1,n,st[top[u]],st[u],k))+=Mod;
        u=fa[top[u]];
    }
    (res+=T.Query(1,1,n,min(st[v],st[u]),max(st[v],st[u]),k))%=Mod;
    return res;
}
int main()
{
    read(n);
    for(register int i=1;i<n;++i)
    {
        int u,v;read(u);read(v);
        insert(u,v);insert(v,u);
    }
    dep[0]=-1;
    init();
    read(q);
    while(q--)
    {
        int u,v,k;read(u);read(v);read(k);
        write(Getans(u,v,k),'\n');
    }
    return 0;
}