2020牛客寒假算法基础集训营1——F.maki和tree【树形DP & 树上DFS】

题目传送门

---

题目描述

有一天，maki拿到了一颗树。所谓树，即没有自环、重边和回路的无向连通图。
这个树有 $n$ 个顶点，$n-1$ 条边。每个顶点被染成了白色或者黑色。
maki想知道，取两个不同的点，它们的简单路径上有且仅有一个黑色点的取法有多少？
注：
①树上两点简单路径指连接两点的最短路。
② $<p,q>$ 和 $<q,p>$ 的取法视为同一种。

---

输入描述:

第一行一个正整数 $n$ 。代表顶点数量。$(1<=n<=100000)$
第二行是一个仅由字符’B’和’W’组成的字符串。第 $i$ 个字符是B代表第 $i$ 个点是黑色，W代表第 $i$ 个点是白色。
接下来的 $n-1$ 行，每行两个正整数 $x$ ，$y$ ，代表 $x$ 点和 $y$ 点有一条边相连 $(1<=x,y<=n)$

---

输出描述:

一个正整数，表示只经过一个黑色点的路径数量。

---

输入

3
WBW
1 2
2 3

---

输出

3

---

说明

树表示如下：

其中只有2号是黑色点。
<1,2>、<2,3>、<1,3>三种取法都只经过一个黑色点。

---

题解

• 官方题解是 $并查集+前缀和$ 维护。并且后台数据没有卡 $暴力dfs每个黑点的$ $N^2$ 解法，（有点水了）
• 对于这个题目，显然DP更佳。
• $定义：dp[i][0]表示子树到i路径上无黑点的个数，dp[i][1]表示有一个黑点的个数$
• $参考cf难度分：1900$

---

AC-Code

#include<bits/stdc++.h>
#define ll long long
using namespace std;

const int maxn = 1e5 + 7;
vector<int>G[maxn << 1];
ll dp[maxn][2]; // dp[i][0]表示子树到i路径上无黑点的个数，dp[i][1]表示有一个黑点的个数
char s[maxn];
ll ans;
void dfs(int u, int fa) {
	if (s[u] == 'B')	dp[u][1] = 1, dp[u][0] = 0;
	else	dp[u][0] = 1, dp[u][1] = 0;
	for (auto v : G[u]) {
		if (v == fa)	continue;
		dfs(v, u);
		if (s[u] == 'B') {
			ans += dp[u][1] * dp[v][0];
			dp[u][1] += dp[v][0];
			dp[u][0] = 0;
		}
		else {
			ans += (dp[u][0] * dp[v][1] + dp[u][1] * dp[v][0]);
			dp[u][1] += dp[v][1];
			dp[u][0] += dp[v][0];
		}
	}
}
int main() {
	int n;
	while (cin >> n) {
		ans = 0;
		memset(dp, 0, sizeof dp);
		for (int i = 0; i <= n; ++i)	G[i].clear();
		cin >> s + 1;
		for (int i = 0; i < n - 1; ++i) {
			int u, v;	cin >> u >> v;
			G[u].push_back(v), G[v].push_back(u);
		}
		dfs(1, 1);
		cout << ans << endl;
	}
}