数据结构(主席树)：COGS 2211. 谈笑风生

2211. 谈笑风生

★★★★ 输入文件：

laugh.in

输出文件：

laugh.out

简单对比
时间限制：3 s 内存限制：512 MB

【问题描述】

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

• 设a和b为T 中的两个不同节点。如果a是b的祖先，那么称“a比b不知道高明到哪里去了”。

• 设a 和 b 为 T 中的两个不同节点。如果 a 与 b 在树上的距离不超过某个给定常数x，那么称“a 与b 谈笑风生”。

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

1. a、b和 c为 T 中三个不同的点，且 a为p 号节点；

2. a和b 都比 c不知道高明到哪里去了；

3. a和b 谈笑风生。这里谈笑风生中的常数为给定的 k。

【输入格式】

输入文件的第一行含有两个正整数n和q，分别代表有根树的点数与询问的个数。接下来n-1行，每行描述一条树上的边。每行含有两个整数u和v,代表在节点u和v之间有一条边。接下来q 行，每行描述一个操作。第i行含有两个整数，分别表示第i个询问的p和k。

【输出格式】输出 q 行，每行对应一个询问，代表询问的答案。

laugh.in

5 3

1 2

1 3

2 4

4 5

2 2

4 1

2 3

laugh.out

3

1

3

好题就要分享，若本题作者认为侵权，请告知我，我会尽快删除。

作者:TenderRun

　　这道题目一看就是要用某种数据结构的。

　　观察它的询问：有两类情况

　　　　①：b点是a点祖先，可以O(1)求出答案。

　　　　②：b是a子树中的一个节点，我们可以用DFS序，那么b点所在子树就是DFS序列中连续的一段，考虑对子树分层，接着用主席树水过了。

#include <iostream>
#include <cstring>
#include <cstdio>
#include <queue>
using namespace std;
const int maxn=300010;
int cnt,fir[maxn],nxt[maxn<<1],to[maxn<<1];
void addedge(int a,int b){
nxt[++cnt]=fir[a];
fir[a]=cnt;
to[cnt]=b;
}
int dep[maxn],sz[maxn],ID[maxn],end[maxn],tot;
void DFS(int node){
sz[node]=1;ID[node]=++tot;
for(int i=fir[node];i;i=nxt[i]){
if(dep[to[i]])continue;
dep[to[i]]=dep[node]+1;
DFS(to[i]);
sz[node]+=sz[to[i]];
}
end[node]=tot;
}
queue<int>q;
int rt[maxn],ch[maxn*30][2],cont;
long long tr[maxn*30];
void Insert(int pre,int &rt,int l,int r,int g){
rt=++cont;
ch[rt][0]=ch[pre][0];
ch[rt][1]=ch[pre][1];
tr[rt]=tr[pre]+sz[g]-1;
if(l==r)return;
int mid=(l+r)>>1;
if(ID[g]<=mid)Insert(ch[pre][0],ch[rt][0],l,mid,g);
else Insert(ch[pre][1],ch[rt][1],mid+1,r,g);
}
long long Query(int pre,int rt,int l,int r,int a,int b){
if(l>=a&&r<=b)return tr[rt]-tr[pre];
int mid=(l+r)>>1;
long long ret=0;
if(mid>=a)ret=Query(ch[pre][0],ch[rt][0],l,mid,a,b);
if(mid<b)ret+=Query(ch[pre][1],ch[rt][1],mid+1,r,a,b);
return ret;
}
int main(){
freopen("laugh.in","r",stdin);
freopen("laugh.out","w",stdout);
int n,Q;
scanf("%d%d",&n,&Q);
for(int i=1,a,b;i<n;i++){
scanf("%d%d",&a,&b);
addedge(a,b);
addedge(b,a);
}
dep[1]=1;
DFS(1);
q.push(1);
int maxd=0;
while(!q.empty()){
int node=q.front();q.pop();maxd=max(maxd,dep[node]);
if(!rt[dep[node]])Insert(rt[dep[node]-1],rt[dep[node]],1,n,node);
else Insert(rt[dep[node]],rt[dep[node]],1,n,node);
for(int i=fir[node];i;i=nxt[i]){
if(dep[to[i]]!=dep[node]+1)continue;
q.push(to[i]);
}
}
int a,k;
long long ans;
while(Q--){
scanf("%d%d",&a,&k);
ans=1ll*(min(k,dep[a]-1))*(sz[a]-1);
ans+=Query(rt[dep[a]],rt[min(dep[a]+k,maxd)],1,n,ID[a],end[a]);
printf("%lld\n",ans);
}
return 0;
}