牛客小白月赛44 E. 变异蛮牛（树形DP）

链接：https://ac.nowcoder.com/acm/contest/11221/E
来源：牛客网

题目描述

幽怨火，憎恨焰，变异蛮牛续执念。

给定一棵根为 11，且是黑点的有根树。

每个白点相邻所有的点都是黑点，每个黑点相邻所有的点都是白点。换句话说，你可以从根结点开始，按照深度对每个点黑白染色。

现在对于一条两个端点分别是 u,vu,v 的链，定义其长度为：包含的黑点个数 −− 包含的白点个数。

请你数一数 长度最大 的链的个数。

输入描述:

全文第一行是 T(1≤T≤105)T(1≤T≤105)，表示数据组数；


接下来 TT 组数据，先输入一行一个正整数表示树的大小 n(1≤n≤2×105,∑n≤3×106)n(1≤n≤2×105,∑n≤3×106)；
接下来输入 n−1n−1 行每行两个正整数 u,v(1≤u,v≤n)u,v(1≤u,v≤n) 表示树的一条边。

输出描述:

输出 TT 行，每行一个整数表示答案。

示例1

输入

复制

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

输出

复制

6

说明

合法的链分别是：{1},{3},{5},{1,2,3},{3,4,5},{1,2,3,4,5}{1},{3},{5},{1,2,3},{3,4,5},{1,2,3,4,5}。

题意等价于问首尾两端点都是黑色的链一共有多少条，可以树dp解决。首先一遍dfs预处理出szb这个数组（szb[x]代表以x为根的子树中的黑色点的数目）。对于每棵子树，如果子树的根x是黑色，那么链分为如下几部分：一个端点是x而另一个黑色端点在x的子树中（对子树答案的贡献是x的每个儿子的sbz值的和）；两个端点分处于x的不同子树中（对子树答案的贡献是\(\Sigma_{i=1}^pszb[i]\times \Sigma_{j=i+1}^pszb[j]\)，这个可以利用前缀和思想在遍历儿子的时候不断累加维护一下）；单独的一个点x（对自述答案的贡献为1）。如果x是白色，那么只有两个端点分处于x的不同子树中这一种情况。

#include <iostream>
#include <vector>
#define int long long
#define N 200005
using namespace std;
int head[N], ver[2 * N], Next[2 * N], c[N], tot, n;
int szb[N];
void add(int x, int y) {
	ver[++tot] = y, Next[tot] = head[x], head[x] = tot;
}
void dfs(int x, int pre, int color) {
	c[x] = color;
	if(color == 1) szb[x]++;
	for(int i = head[x]; i; i = Next[i]) {
		int y = ver[i];
		if(y == pre) continue;
		dfs(y, x, color ^ 1);
		szb[x] += szb[y];
	}
}
int ans[N], totans = 0;
void dfs2(int x, int pre) {
	int tmp = 0;
	bool flag = 0;
	for(int i = head[x]; i; i = Next[i]) {
		int y = ver[i];
		if(y == pre) continue;
		dfs2(y, x);
		ans[x] += szb[y] * tmp;
		tmp += szb[y];
	}
	if(c[x] == 1) ans[x] += 1 + tmp;
	totans += ans[x];
}
signed main() {
	int t;
	cin  >> t;
	while(t--) {
		tot = 0;
		cin >> n;
		for(int i = 1; i <= n; i++) {
			head[i] = 0;
			totans =0;
			szb[i] = 0;
			ans[i] = 0;
		}
		for(int i = 1; i <= n - 1; i++) {
			int u, v;
			cin >> u >> v;
			add(u, v);
			add(v, u);
		}
		dfs(1, 0, 1);
		dfs2(1, 0);
		cout << totans << endl;
	}
	return 0;
}