【bzoj3653】谈笑风生 DFS序+树状数组

题目描述

给出一棵以1为根的有根树，q次询问，每次询问给出a和k，求点对 (b,c) 的数目，满足：a、b、c互不相同，b与a距离不超过k，且a和b都是c的祖先。

输入

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

输出

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

样例输入

5 3
1 2
1 3
2 4
4 5
2 2
4 1
2 3

样例输出

3
1
3

---

题解

DFS序+树状数组

考虑我们所求的是什么：

显然a和b一定具有祖先关系，否则不存在满足a和b均为祖先的点c。

当b是a的祖先时，b在a到根的路径上，c在a的子树中，可以直接将答案计算出来。

当a是b的祖先时，b在a的子树内且与a的距离不超过k，c在b的子树中。

“b在a的子树内”可以看成是DFS序中的位置限制关系，“与a距离不超过k”可以看成是深度限制关系。如果我们把子树大小-1看成一个点的权值（c能取的位置数），那么要求出的就是DFS序位置在某区间，深度小于等于某值的点的权值和。

将 (DFS序位置，深度） 看作二维平面上的点，要求的就是矩形内的点的权值和。离线处理，把一个询问拆成两个，按照DFS序位置排序，扫一遍，并使用树状数组维护前缀和即可。

时间复杂度$O(n\log n)$

#include <cstdio>
#include <algorithm>
#define N 300010
using namespace std;
typedef long long ll;
struct data
{
	int c , d , opt , id;
	bool operator<(const data &a)const {return c < a.c;}
}a[N << 1];
int n , head[N] , to[N << 1] , next[N << 1] , cnt , deep[N] , si[N] , pos[N] , ref[N] , tot;
ll f[N] , ans[N];
inline void add(int x , int y)
{
	to[++cnt] = y , next[cnt] = head[x] , head[x] = cnt;
}
void dfs(int x , int fa)
{
	int i;
	deep[x] = deep[fa] + 1 , si[x] = 1 , pos[x] = ++tot , ref[tot] = x;
	for(i = head[x] ; i ; i = next[i])
		if(to[i] != fa)
			dfs(to[i] , x) , si[x] += si[to[i]];
}
inline void update(int x , int a)
{
	int i;
	for(i = x ; i <= n ; i += i & -i) f[i] += a;
}
inline ll query(int x)
{
	int i;
	ll ans = 0;
	for(i = x ; i ; i -= i & -i) ans += f[i];
	return ans;
}
int main()
{
	int m , i , x , y , p = 0;
	scanf("%d%d" , &n , &m);
	for(i = 1 ; i < n ; i ++ ) scanf("%d%d" , &x , &y) , add(x , y) , add(y , x);
	dfs(1 , 0);
	for(i = 1 ; i <= m ; i ++ )
	{
		scanf("%d%d" , &x , &y) , ans[i] = (ll)(si[x] - 1) * min(y , deep[x] - 1);
		a[i].c = pos[x] , a[i].opt = -1 , a[i + m].c = pos[x] + si[x] - 1 , a[i + m].opt = 1;
		a[i].d = a[i + m].d = min(deep[x] + y , n) , a[i].id = a[i + m].id = i;
	}
	sort(a + 1 , a + m * 2 + 1);
	for(i = 1 ; i <= m * 2 ; i ++ )
	{
		while(p <= n && p < a[i].c) p ++ , update(deep[ref[p]] , si[ref[p]] - 1);
		ans[a[i].id] += a[i].opt * query(a[i].d);
	}
	for(i = 1 ; i <= m ; i ++ ) printf("%lld\n" , ans[i]);
	return 0;
}