cogs2478 简单的最近公共祖先

【题目描述】

给定一棵有n个节点的有根树，根节点为1，每个节点有一个权值wi，求

即求所有无序节点对的LCA的权值之和。

树的节点编号为1~n，LCA表示两节点的最近公共祖先，即在它们的所有公共祖先中离根节点最远的节点。

【输入格式】

第一行一个整数n，表示节点数。

第二行n个正整数，表示每个点的权值。

以下n-1行每行两个整数x,y，表示树上有一条边连接节点x和节点y。

【输出格式】

一个整数，表示答案。

【样例输入】

3
1 2 3
1 2
1 3

【样例输出】

9

【数据范围与约定】

对于30%的数据，n<=1000。

对于60%的数据，n<=100000。

对于100%的数据，1<=n<=1000000，0<wi<=1000000。

呃。。。很明显不是枚举所有点对。。。

容易发现，对于一个节点来说，它作为lca的时候当且仅当任意两个子节点属于该节点的两个不相同的子树。

然后就乘法原理＋前缀和优化就行了。

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <vector>
using namespace std;
#define FILENAME "easy_LCA"

typedef long long ll;

const int N = 2000010;

int n, head[N], rest[N], to[N], w[N], size[N], tot;

void add(int u, int v) {
	tot ++;
	to[tot] = v;
	rest[tot] = head[u];
	head[u] = tot;
}

ll ret;

void dfs(int u) {
	size[u] = 1;
	for(int i = head[u] ; i ; i = rest[i]) {
		if(size[to[i]] == 0) {
			dfs(to[i]);
			ret += (ll)size[u] * (ll)size[to[i]] * (ll)w[u];
			size[u] += size[to[i]];
		}
	}
	ret += w[u];
}


int main() {
	freopen(FILENAME ".in", "r", stdin);
	freopen(FILENAME ".out", "w", stdout);
	scanf("%d", &n);
	for(int i = 1 ; i <= n ; i ++) {
		scanf("%d", &w[i]);
	}
	for(int i = 1, u, v ; i < n ; i ++) {
		scanf("%d%d", &u, &v);
		add(u, v);
		add(v, u);
	}
	dfs(1);
	printf("%lld\n", ret);
}