[BZOJ5293]BJOI2018：求和

BZOJ5293

你都做BJOI了你会做不来这道题？

Code：

#include
#define mod 998244353
#define ll long long
#define int long long
using namespace std;
inline int read(){
	int res=0,f=1;char ch=getchar();
	while(!isdigit(ch)) {if(ch=='-') f=-f;ch=getchar();}
	while(isdigit(ch)) {res=(res<<1)+(res<<3)+(ch^48);ch=getchar();}
	return res*f;
}
const int N=3e5+5,K=51;
int vis[N<<1],head[N<<1],nxt[N<<1],c[N<<1],tot=0;
inline void add(int x,int y,int z){vis[++tot]=y,nxt[tot]=head[x];head[x]=tot;c[tot]=z;}
int dis[N][K],dep[N],fa[N],siz[N],hson[N],pt[N];
inline int ksm(int x,int y){int res=1;for(;y;y>>=1,x=x*x%mod) if(y&1) res=res*x%mod; return res;}
void dfs1(int v){
	pt[v]=siz[v]=1;
	for(int i=head[v];i;i=nxt[i]){
		int y=vis[i];
		if(pt[y]) continue;
		fa[y]=v;dep[y]=dep[v]+1;
		for(int j=1;j<K;j++)
		dis[y][j]=((ll)dis[v][j]+ksm(dep[y],j))%mod;
		dfs1(y);
		siz[v]+=siz[y];
		if(siz[y]>siz[hson[v]]) hson[v]=y;
	}
}
int top[N];
void dfs2(int v){
	if(v==hson[fa[v]]) top[v]=top[fa[v]];
	else top[v]=v;	for(int i=head[v];i;i=nxt[i])
		if(!top[vis[i]]) dfs2(vis[i]);
}
inline int lca(int x,int y){
	while(top[x]!=top[y])
		dep[top[x]]>dep[top[y]]?x=fa[top[x]]:y=fa[top[y]];
	return dep[x]<dep[y]?x:y;
}
inline int d(int x,int y,int k){return ((dis[x][k]+dis[y][k]-dis[lca(x,y)][k]+mod-dis[fa[lca(x,y)]][k]+mod)+mod)%mod;}
signed main(){
	int n=read();
	for(int x,y,i=1;i<n;i++) x=read(),y=read(),add(x,y,1),add(y,x,1);
	dfs1(1);top[1]=1;dfs2(1);
	int m=read();
	for(int x,y,z,i=1;i<=m;i++) x=read(),y=read(),z=read(),cout<<d(x,y,z)<<"\n";
	return 0;
}