loj #3088 「GXOI / GZOI2019」旧词 离线+树链剖分

Description

---

n个节点的树，m次询问(x,y)
求$$\sum _{i=1}^x{\left(dep[lca\left(i,y\right)]\right)}^k$$
其中k是一个给定的常数

Solution

---

观察k=1的时候要怎么做。我们离线按x排序，对于一个节点t把根到t路径上的所有点全部+1。那么y和前面所有点的lca的深度之和就是y到根路径上所有点权的和

拓展一下就可以发现，我们只需要把+1变成加上dep[x]^k-dep[fa[x]] ^k就可以了

Code

---

#include <stdio.h>
#include <string.h>
#include <algorithm>

#define rep(i,st,ed) for (int i=st;i<=ed;++i)
#define lson now<<1
#define rson now<<1|1

typedef long long LL;
const int MOD=998244353;
const int N=500005;

struct edge {int y,next;} e[N*2];
struct Q {int x,y,id;} q[N];

LL sum[N<<2],val[N<<2],tag[N<<2],r[N],ans[N];
int pos[N],dfn[N],size[N],fa[N],dep[N],bl[N];
int ls[N],edCnt;
int p[N];

bool np[N];

int read() {
	int x=0,v=1; char ch=getchar();
	for (;ch<'0'||ch>'9';v=(ch=='-')?(-1):v,ch=getchar());
	for (;ch<='9'&&ch>='0';x=x*10+ch-'0',ch=getchar());
	return x*v;
}

void upd(LL &x,LL v) {
	x+=v,(x>=MOD)?(x-=MOD):0;
}

void add_edge(int x,int y) {
	e[++edCnt]=(edge) {y,ls[x]}; ls[x]=edCnt;
}

LL ksm(LL x,LL dep,LL res=1) {
	for (x%=MOD;dep;dep>>=1,x=x*x%MOD) {
		(dep&1)?(res=res*x%MOD):0;
	} return res;
}

void pre(int n,int k) {
	r[1]=1;
	rep(i,2,n) {
		if (!np[i]) p[++p[0]]=i,r[i]=ksm(i,k);
		for (int j=1;j<=p[0]&&i*p[j]<=n;++j) {
			np[i*p[j]]=1;
			r[i*p[j]]=r[i]*r[p[j]]%MOD;
			if (i%p[j]==0) break;
		}
	}
}

void dfs1(int x) {
	size[x]=1;
	for (int i=ls[x];i;i=e[i].next) {
		dep[e[i].y]=dep[x]+1;
		dfs1(e[i].y),size[x]+=size[e[i].y];
	}
}

void dfs2(int x,int up) {
	bl[x]=up,pos[x]=++pos[0];
	int mx=0; dfn[pos[x]]=x;
	for (int i=ls[x];i;i=e[i].next) {
		if (size[e[i].y]>size[mx]) mx=e[i].y;
	}
	if (!mx) return ;
	dfs2(mx,up);
	for (int i=ls[x];i;i=e[i].next) {
		if (e[i].y!=mx) dfs2(e[i].y,e[i].y);
	}
}

void push_down(int now,int tl,int tr) {
	if (!tag[now]) return ;
	LL w=tag[now]; tag[now]=0;
	upd(tag[lson],w),upd(tag[rson],w);
	upd(sum[lson],w*val[lson]%MOD);
	upd(sum[rson],w*val[rson]%MOD);
}

void modify(int now,int tl,int tr,int l,int r) {
	if (tl>=l&&tr<=r) {
		upd(tag[now],1),upd(sum[now],val[now]%MOD);
		return ;
	}
	int mid=(tl+tr)>>1;
	push_down(now,tl,tr);
	if (l<=mid) modify(lson,tl,mid,l,r);
	if (mid+1<=r) modify(rson,mid+1,tr,l,r);
	sum[now]=sum[lson],upd(sum[now],sum[rson]);
}

LL query(int now,int tl,int tr,int l,int r) {
	if (tl>=l&&tr<=r) return sum[now];
	int mid=(tl+tr)>>1; LL res=0;
	push_down(now,tl,tr);
	if (l<=mid) upd(res,query(lson,tl,mid,l,r));
	if (mid+1<=r) upd(res,query(rson,mid+1,tr,l,r));
	sum[now]=sum[lson],upd(sum[now],sum[rson]);
	return res;
}

void build(int now,int tl,int tr) {
	if (tl==tr) {
		int x=dfn[tl];
		val[now]=r[dep[x]]-r[dep[fa[x]]];
		upd(val[now],MOD);
		return ;
	}
	int mid=(tl+tr)>>1;
	build(lson,tl,mid),build(rson,mid+1,tr);
	val[now]=val[lson],upd(val[now],val[rson]);
}

void change(int x,int y) {
	for (;bl[x]!=bl[y];x=fa[bl[x]]) {
		if (dep[bl[x]]<dep[bl[y]]) std:: swap(x,y);
		modify(1,1,pos[0],pos[bl[x]],pos[x]);
	}
	if (pos[x]>pos[y]) std:: swap(x,y);
	modify(1,1,pos[0],pos[x],pos[y]);
}

LL ask(int x,int y) {
	LL res=0,tmp;
	for (;bl[x]!=bl[y];x=fa[bl[x]]) {
		if (dep[bl[x]]<dep[bl[y]]) std:: swap(x,y);
		tmp=query(1,1,pos[0],pos[bl[x]],pos[x]);
		upd(res,tmp);
	}
	if (pos[x]>pos[y]) std:: swap(x,y);
	tmp=query(1,1,pos[0],pos[x],pos[y]);
	upd(res,tmp);
	return res;
}

bool cmp(Q a,Q b) {
	return a.x<b.x;
}

int main(void) {
	freopen("poetry.in","r",stdin);
	freopen("poetry.out","w",stdout);
	int n=read(),m=read(),k=read();
	pre(n,k);
	rep(i,2,n) {
		fa[i]=read();
		add_edge(fa[i],i);
	}
	dfs1(dep[1]=1);
	dfs2(1,1);
	build(1,1,pos[0]);
	rep(i,1,m) {
		q[i].x=read(),q[i].y=read();
		q[i].id=i;
	}
	std:: sort(q+1,q+m+1,cmp);
	for (int i=1,j=1;i<=m;++i) {
		while (j<=q[i].x) {
			change(1,j); j++;
		}
		ans[q[i].id]=ask(1,q[i].y);
	}
	rep(i,1,m) printf("%lld\n", ans[i]);
	return 0;
}