牛客小白月赛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;
}