【bzoj2783】[JLOI2012]树 树上倍增

题目描述

在这个问题中，给定一个值S和一棵树。在树的每个节点有一个正整数，问有多少条路径的节点总和达到S。路径中节点的深度必须是升序的。节点1是根节点，根的深度是0，它的儿子节点的深度为1。路径不必一定从根节点开始。

输入

第一行是两个整数N和S，其中N是树的节点数。

第二行是N个正整数，第i个整数表示节点i的正整数。

接下来的N-1行每行是2个整数x和y，表示y是x的儿子。

输出

输出路径节点总和为S的路径数量。

样例输入

3 3
1 2 3
1 2
1 3

样例输出

2

---

题解

树上倍增

O(n)的做法太高端了，于是我选择了O(nlogn)

注意到路径中节点深度必须是升序的，所以路径一定是从某个点向根节点的道路。

先dfs对每个点进行预处理，处理出每个点的2ⁱ祖先和自己到2ⁱ祖先的点权之和（不包括自己）。

然后枚举每个点寻找答案，若总和小于s，则将该点上移继续查找（类似倍增LCA）；若等于s则加入到答案中。

路径可能只包括一个点，注意判断。

#include <cstdio>
#include <algorithm>
#define N 100010
using namespace std;
int a[N] , head[N] , to[N] , next[N] , cnt , log[N] , fa[N][20] , sum[N][20] , deep[N];
void add(int x , int y)
{
	to[++cnt] = y , next[cnt] = head[x] , head[x] = cnt;
}
void dfs(int x)
{
	int i;
	for(i = 1 ; i <= log[deep[x]] ; i ++ )
		fa[x][i] = fa[fa[x][i - 1]][i - 1] , sum[x][i] = sum[x][i - 1] + sum[fa[x][i - 1]][i - 1];
	for(i = head[x] ; i ; i = next[i])
		fa[to[i]][0] = x , sum[to[i]][0] = a[x] , deep[to[i]] = deep[x] + 1 , dfs(to[i]);
}
int main()
{
	int n , m , i , j , k , p , x , y , ans = 0;
	scanf("%d%d" , &n , &m);
	for(i = 1 ; i <= n ; i ++ ) scanf("%d" , &a[i]);
	for(i = 1 ; i < n ; i ++ ) scanf("%d%d" , &x , &y) , add(x , y);
	for(i = 2 ; i <= n ; i ++ ) log[i] = log[i >> 1] + 1;
	dfs(1);
	for(i = 1 ; i <= n ; i ++ )
	{
		if(a[i] == m) ans ++ ;
		for(k = a[i] , p = i , j = log[deep[i]] ; ~j ; j -- )
		{
			if(deep[p] >= (1 << j))
			{
				if(sum[p][j] + k < m) k += sum[p][j] , p = fa[p][j];
				else if(sum[p][j] + k == m) ans ++ ;
			}
		}
	}
	printf("%d\n" , ans);
	return 0;
}