题解 [Codeforces1156D] 0-1-Tree

题解 [Codeforces1156D] 0-1-Tree

题目来源：Codeforces Educational Round #64

tag：树形 DP，计数

---

Solution

很容易想到树形 DP，设 $f_{x,0/1/2/3}$ 表示 $x$ 的子树中有多少个点，到 $x$ 的路径为全 $\text{0}$、全 $\text{1}$、先 $\text{0}$ 再 $\text{1}$，先 $\text{1}$ 再 $\text{0}$。注意四个状态互相没有重合，且单点不算在路径里。

设当前考虑点 $x$ 的答案，每遍历到 $x$ 的一棵子树 $y$，就要先计算 $y$ 和已经遍历过的 $x$ 的子树的答案，再来更新 $f_x$。否则会存在类似于 $y\to x\to y$ 的路径，不符合题意。所以先统计答案，再考虑状态转移。

统计答案时，先根据 $w(x,y)$ 分类。当 $w(x,y)=1$ 时，路径可以分为三类，其中前两类不考虑全 $\text{1}$，第三类为全 $\text{1}$：

1. $y\to x$，$y$ 的子树中可以取 $f_{y,0}+f_{y,2}$ 个点，$x$ 已经遍历过的子树中可以取 $f_{x,1}+1$ 个点。（因为 $x$ 本身也可以取）
2. $x\to y$，$x$ 已经遍历过的子树中可以取 $f_{x,0}+f_{x,2}$ 个点，$y$ 的子树中可以取 $f_{y,1}+1$ 个点。（因为 $y$ 本身也可以取）
3. 全 $\text{1}$ 路径，$x$ 已经遍历过的子树中可以取 $f_{x,1}+1$ 个点，$y$ 的子树中可以取 $f_{y+1}$ 个点。因为方向有两种，所以答案要再乘上 $2$。

所以答案累加上 $(f_{y,0}+f_{y,2})(f_{x,1}+1)+(f_{x,0}+f_{x,2})(f_{y,1}+1)+2(f_{x,1}+1)(f_{y,1}+1)$。

状态转移式即为：

$$\{\begin{array}{l}{f}_{x,1}=\sum _{y\in Son(x),w(x,y)=1}{f}_{y,1}+1\\ f_{x,2}=\sum _{y\in Son(x),w(x,y)=1}{f}_{y,0}+f_{y,2}\end{array}$$

同理，当 $w(x.y)=0$ 时：

答案累加上 $(f_{y,1}+f_{y,3})(f_{x,0}+1)+(f_{x,1}+f_{x,3})(f_{y,0}+1)+2(f_{x,0}+1)(f_{y,0}+1)$。

状态转移：

$$\{\begin{array}{l}{f}_{x,0}=\sum _{y\in Son(x),w(x,y)=0}{f}_{y,0}+1\\ f_{x,3}=\sum _{y\in Son(x),w(x,y)=0}{f}_{y,1}+f_{y,3}\end{array}$$

const int N = 2e5 + 10;
long long ans, f[N][4];
vector<pii> G[N];
int n;

void dfs(int x, int fa){
	for(int i = 0; i < G[x].size(); ++ i){
		int y = G[x][i].first, z = G[x][i].second;
		if(y == fa){
			continue;
		}
		dfs(y, x);
		if(z){
			ans += 2 * (f[y][1] + 1) * (f[x][1] + 1);
			ans += (f[y][0] + f[y][2]) * (f[x][1] + 1);
			ans += (f[x][0] + f[x][2]) * (f[y][1] + 1);
			f[x][1] += f[y][1] + 1;
			f[x][2] += f[y][2] + f[y][0];
		} else {
			ans += 2 * (f[y][0] + 1) * (f[x][0] + 1);
			ans += (f[y][1] + f[y][3]) * (f[x][0] + 1);
			ans += (f[x][1] + f[x][3]) * (f[y][0] + 1);
			f[x][0] += f[y][0] + 1;
			f[x][3] += f[y][3] + f[y][1];
		}
	}
}

void solve(){
	n = rdi;
	for(int i = 1; i < n; ++ i){
		int x = rdi;
		int y = rdi;
		int z = rdi;
		G[x].push_back(make_pair(y, z));
		G[y].push_back(make_pair(x, z));
	}
	dfs(1, 0);
	writen(ans);
}