Codeforces 1156D 0-1-Tree
题目链接
https://codeforces.com/contest/1156/problem/D
题目大意
给你一个 $n$ 个节点的树，每条边上有权值 $0$ 或 $1$ 。问经过权值为 $1$ 的边之后不通过权值为 $0$ 的有向路径共有多少种。
题目分析
很明显这是一道 $DP$ 题。考虑定义 $f[x][0/1]$ 为从子树出发走到这个点且最后一个边为 $0/1$ 的起始点共有多少个， $g[x][0/1]$ 为从这个点出发，下一条边为 $0/1$，走到子树中共有多少个终点。那么考虑到经过 $1$ 后不能经过 $0$ ，那么转移的时候就不能往这个方向转移即可。计算时相乘即可。
代码
#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cstring>
#include <string>
#include <algorithm>
#include <vector>
using namespace std;
const int maxn = 2e5 + 5;
int n, m, t;
vector <int> e[maxn];
vector <int> w[maxn];
long long f[maxn][2];
long long g[maxn][2];
long long ans = 0;
int read(){
int x;
scanf("%d", &x);
return x;
}
int dfs(int x, int fa){
for(int i=0;i<e[x].size();i++){
int y = e[x][i];
int c = w[x][i];
if(y != fa){
dfs(y, x);
if(c){
f[x][1] += f[y][0] + f[y][1] + 1;
g[x][1] += g[y][1] + 1;
}else{
f[x][0] += f[y][0] + 1;
g[x][0] += g[y][0] + g[y][1] + 1;
}
}
}
ans += f[x][0] + f[x][1] + g[x][0] + g[x][1];
for(int i=0;i<e[x].size();i++){
int y = e[x][i];
int c = w[x][i];
if(y != fa){
if(c){
ans += (f[y][0] + f[y][1] + 1) * (g[x][1] - (g[y][1] + 1));
}else{
ans += (f[y][0] + 1) * (g[x][0] - (g[y][0] + g[y][1] + 1) + g[x][1]);
}
}
}
return 0;
}
int main(){
int i, j;
int x, y, c;
n = read();
for(i=1;i<n;i++){
x = read();
y = read();
c = read();
e[x].push_back(y);
e[y].push_back(x);
w[x].push_back(c);
w[y].push_back(c);
}
dfs(1, -1);
printf("%lld\n", ans);
return 0;
}