Solution
感觉很久没有自己调出来一道 (DP) 了。
我们分类 (DP)。令 (f[i][0]) 为以这个点为根的子树内部起点为 (i) 的类似 (00011111) 与 (00000000) 串的个数, (f[i][1]) 为以这个点为根的子树内部起点为 (i) 的类似 (11111111) 串的个数。
首先我们讨论 (x),(y),其中 (y) 是 (x) 的爸爸。
考虑计算 (f) 数组。易得无论 (w(x,y)) 是什么值,首先爸爸的值是儿子的相应边权的值加一。考虑 (w(x,y)==0) 的情况,还要加上 (f[x][1])。(注意还要加上连接 (x) 与 (y) 的边)
好了我们令 (g[i][0]) 为以这个点为根起点为 (i) 的类似 (00011111) 与 (00000000) 串的个数, (g[i][1]) 为以这个点为根,起点为 (i) 的类似 (11111111) 串的个数,(g) 的初始值赋为 (f)。(代码实现里我是直接用 (f),但是推导过程是有区别的)
好了我们放图(画质感人):
分类讨论:
- (w(x,y)==1):儿子不可能继承父亲的 (0) 情况,所以直接用爸爸减儿子得到其他儿子的值再把儿子加回去。(ans=g[y][1]-f[x][1]+f[x][1]+f[x][0])。
- (w(x,y)==0):这里有个坑。父亲在算 (x) 的时候,不仅会加 (f[x][0]),还会加 (f[x][1]),方法同上,最后还要记得加上父亲的全 (1) 串个数。
Code
#include <cstdio>
typedef long long ll;
const int N = 2e5 + 5;
int n, head[N], nxt[N << 1], to[N << 1], w[N << 1], cnt;
ll f[N][2], ans;
int read() {
int x = 0, f = 1; char s;
while((s = getchar()) > '9' || s < '0') if(s == '-') f = -1;
while(s <= '9' && s >= '0') x = (x << 1) + (x << 3) + (s ^ 48), s = getchar();
return x * f;
}
void addEdge(const int u, const int v, const int val) {
nxt[++ cnt] = head[u], to[cnt] = v, head[u] = cnt, w[cnt] = val;
}
void dfs1(const int u, const int fa) {
for(int i = head[u]; i; i = nxt[i]) {
int v = to[i];
if(v == fa) continue;
dfs1(v, u);
f[u][w[i]] += f[v][w[i]] + 1;
if(! w[i]) f[u][0] += f[v][1];
}
}
void dfs2(const int u, const int fa) {
for(int i = head[u]; i; i = nxt[i]) {
int v = to[i];
if(v == fa) continue;
if(w[i] == 1) f[v][1] = f[u][1];
if(! w[i]) f[v][0] = f[u][0] - f[v][1] + f[u][1];
dfs2(v, u);
}
}
int main() {
int u, v, w;
n = read();
for(int i = 1; i < n; ++ i) {
u = read(), v = read(), w = read();
addEdge(u, v, w); addEdge(v, u, w);
}
dfs1(1, 0); dfs2(1, 0);
for(int i = 1; i <= n; ++ i) ans += f[i][0] + f[i][1];
printf("%lld
", ans);
return 0;
}