BZOJ 5293 [Bjoi2018]求和

题目大意

多次询问树上一条路径上的点的深度的$k$次方。

虽然每次的$k$会变但是$k \le 50$还说啥了……直接暴力维护50个权值，树链剖分就可以了。

#include <cstdio>
#include <cctype>
#include <algorithm>
using namespace std;
#define mo 998244353
int n,m,lg[300010],fa[300010][25];
long long dep[300010][51];
inline char nc()
{
	static char buf[100000],*p1,*p2;
	return p1==p2&&(p2=(p1=buf)+fread(buf,1,100000,stdin),p1==p2)?EOF:*p1++;
}
inline int ri()
{
	int ret=0;char tmp=nc();
	while(!isdigit(tmp)) tmp=nc();
	while(isdigit(tmp)) ret=ret*10+(tmp-'0'),tmp=nc();
	return ret;
}
int to[600010],nxt[600010],head[300010],tot;
inline void ae(int x,int y){to[++tot]=y,nxt[tot]=head[x],head[x]=tot;}
void dfs1(int pos)
{
	int i;
	dep[pos][1]=dep[fa[pos][0]][1]+1;
	for(i=1;i<=lg[n];++i) fa[pos][i]=fa[fa[pos][i-1]][i-1];
	for(i=head[pos];i;i=nxt[i]) if(to[i]!=fa[pos][0]) fa[to[i]][0]=pos,dfs1(to[i]);
}
void dfs2(int pos)
{
	int i;
	for(i=1;i<=50;++i) dep[pos][i]=(dep[pos][i]+dep[fa[pos][0]][i])%mo;
	for(i=head[pos];i;i=nxt[i]) if(to[i]!=fa[pos][0]) dfs2(to[i]);
}
int lca(int x,int y)
{
	int i;
	if(dep[x][1]>dep[y][1]) swap(x,y);
	for(i=lg[n];~i;--i) if(dep[fa[y][i]][1]>=dep[x][1]) y=fa[y][i];
	for(i=lg[n];~i;--i) if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
	return x==y?x:fa[x][0];
}
int main()
{
	int i,j,t,x,y,z;
	n=ri();
	for(i=2;i<=n;++i) x=ri(),y=ri(),ae(x,y),ae(y,x),lg[i]=lg[i>>1]+1;
	dep[0][1]=-1,dfs1(1),dep[0][1]=0;
	for(i=2;i<=50;++i) for(j=1;j<=n;++j) dep[j][i]=dep[j][i-1]*dep[j][1]%mo;
	dfs2(1);
	dep[0][1]=-1;
	m=ri();
	while(m--)
	{
		x=ri(),y=ri(),z=ri(),t=lca(x,y);
		printf("%lld\n",((dep[x][z]+dep[y][z]-dep[t][z]-(t==1?0:dep[fa[t][0]][z]))%mo+mo)%mo);
	}
	return 0;
}