CodeForces - 1156D 0-1-Tree

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\n", ans);
    return 0;
}