luogu P3899 [湖南集训]更为厉害 |线段树合并

题目描述

设 \(\text T\) 为一棵有根树，我们做如下的定义：

设 \(a\) 和 \(b\) 为 \(\text T\) 中的两个不同节点。如果 \(a\) 是 \(b\) 的祖先，那么称“\(a\) 比 \(b\) 更为厉害”。

设 \(a\) 和 \(b\) 为 \(\text T\) 中的两个不同节点。如果 \(a\) 与 \(b\) 在树上的距离不超过某个给定常数 \(x\)，那么称“ \(a\) 与 \(b\) 彼此彼此”。

给定一棵 \(n\) 个节点的有根树 \(\text T\)，节点的编号为 \(1\) 到 \(n\)，根节点为 \(1\) 号节点。
你需要回答 \(q\) 个询问，询问给定两个整数\(p\) 和 \(k\)，问有多少个有序三元组 \((a,b,c)\) 满足：

\(a,b,c\) 为 \(\text T\) 中三个不同的点，且 \(a\) 为 \(p\) 号节点；

\(a\) 和 \(b\) 都比 \(c\) 更为厉害；

\(a\) 和 \(b\) 彼此彼此。这里彼此彼此中的常数为给定的 \(k\)。

输入格式

输入文件的第一行含有两个正整数 \(n\) 和 \(q\)，分别代表有根树的点数与询问的个数。

接下来 \(n − 1\) 行，每行描述一条树上的边。每行含有两个整数 \(u\) 和 \(v\)，代表在节点 \(u\) 和 \(v\) 之间有一条边。

接下来 \(q\) 行，每行描述一个操作。第 \(i\) 行含有两个整数，分别表示第 \(i\) 个询问的 \(p\) 和 \(k\)。

输出格式

输出 \(q\) 行，每行对应一个询问，代表询问的答案。

---

调了我好久

---

#include <bits/stdc++.h>
using namespace std;
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define int long long
inline int read(){
    int x=0; char ch=getchar();
    while(ch<'0'||ch>'9')ch=getchar();
    while(ch>='0'&&ch<='9'){ x=(x<<1)+(x<<3)+(ch^48); ch=getchar(); }
    return x;
}
const int N=5e5+5,M=1e7+1;
int nxt[N<<1],head[N],go[N<<1],tot;
inline void add(int u,int v){
	nxt[++tot]=head[u],head[u]=tot,go[tot]=v;
	nxt[++tot]=head[v],head[v]=tot,go[tot]=u;
}
vector<pair<int,int> >Q[N];
int val[M],sum[M],ls[M],rs[M],root[N],cnt;
#define mid ((l+r)>>1)
inline void pushup(int p){
	val[p]=val[ls[p]]+val[rs[p]];
	sum[p]=sum[ls[p]]+sum[rs[p]];
}
void update(int &p,int l,int r,int pos,int d){
	if(!p)p=++cnt;
	if(l==r){ val[p]+=d; sum[p]++; return; }
	if(pos<=mid)update(ls[p],l,mid,pos,d);
	else update(rs[p],mid+1,r,pos,d);
	pushup(p);
}
int Val,Sum;
void query(int p,int l,int r,int L,int R){
	if(L<=l&&r<=R){ Val+=val[p],Sum+=sum[p]; return; }
	if(L<=mid)query(ls[p],l,mid,L,R);
	if(R>mid)query(rs[p],mid+1,r,L,R);
}
int merge(int u,int v,int l,int r){
	if(!u||!v)return u|v;
	int rt=++cnt;
	val[rt]=val[u]+val[v],sum[rt]=sum[u]+sum[v];
	if(l==r)return rt;
	ls[rt]=merge(ls[u],ls[v],l,mid);
	rs[rt]=merge(rs[u],rs[v],mid+1,r);
	return rt;
}
int dep[N],Ans[N],siz[N],n,q;
void dfs(int u,int fa){
	dep[u]=dep[fa]+1,siz[u]=1;
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(v==fa)continue;
		dfs(v,u);
		siz[u]+=siz[v];
	}
	for(int i=0;i<Q[u].size();i++){
		int k=Q[u][i].fi;
		int ans= min(dep[u]-1,k) * (siz[u]-1) ;
		Val=0,Sum=0;
		query(root[u],1,n,dep[u],k+dep[u]); 
		ans+= Val- Sum*dep[u] -Sum ;
		ans+= (sum[root[u]]-Sum)*k;
		Ans[Q[u][i].se]=ans;
	}
	update(root[u],1,n,dep[u],dep[u]);
	if(fa)root[fa]=merge(root[fa],root[u],1,n);
}

signed main(){
	n=read(),q=read();
	for(int i=1;i<n;i++)add(read(),read());
	for(int i=1,x,y;i<=q;i++){
		x=read(),y=read();
		Q[x].pb(mp(y,i));
	}
	dfs(1,0);
	for(int i=1;i<=q;i++)printf("%lld\n",Ans[i]);
}