【概率期望DP 换根】luogu_P4284 [SHOI2014]概率充电器

题意

概率充电器由\(n-1\)条导线连通了\(n\)个充电元件。
进行充电时，每条导线是否可以导电以概率决定，每一个充电元件自身是否直接进行充电也由概率决定。
随后电能可以从直接充电的元件经过通电的导线使得其他充电元件进行间接充电。
求进入充电状态的元件个数的期望？

对于\(100\%\)的数据，\(n≤500000，0≤p,qi≤100\)。

思路

可以发现给出的是一棵树，自然想到树形dp。
对于一个元件，它可以自身直接充电，也可以从子树里间接充电，又或是子树外间接充电。
设\(f_i\)为对于自己的子树，第\(i\)个元件充不到电的概率。
根据对立事件\(f_i=(1-p_i)\prod_{y\in son_i}1-e_i*(1-f_y)\)。
那么还要考虑子树外的贡献，考虑换根。
设\(g_i\)为\(i\)节点充不到电的概率，对于根\(x\to y\)，要将\(y\)对\(x\)的贡献减去，再用\(x\)的答案去贡献\(y\)。

代码

#include <cstdio>
#include <iostream>
#include <algorithm>
#define getchar() (p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 1 << 21, stdin), p1 == p2) ? EOF : *p1++)

char buf[1 << 23], *p1 = buf, *p2 = buf, obuf[1 << 23], *O = obuf;

inline int read() {
	int res = 0, f = 1;
	char c = getchar();
	while (!isdigit(c)) {
		if (c == '-') f = -1;
		c = getchar();
	}
	while (isdigit(c))
		res = res * 10 + (c ^ 48), c = getchar();
	return res * f;
}

int n, tot;
double ans;
int ver[1000001], next[1000001], head[500001];
double edge[1000001], f[500001], p[500001], g[500001];

void add(int u, int v, double w) {
	ver[++tot] = v;
	next[tot] = head[u];
	edge[tot] = w;
	head[u] = tot;
}

void dfs1(int fa, int now) {
	f[now] = 1 - p[now];
	for (int i = head[now], y; i; i = next[i])
		if (fa != (y = ver[i])) {
			dfs1(now, y);
			f[now] *= 1 - edge[i] * (1 - f[y]);
		}
}

void dfs2(int fa, int now) {//换根
	ans += 1 - g[now];
	for (int i = head[now], y; i; i = next[i])
		if (fa != (y = ver[i])) {
			if (edge[i] * (1 - f[y]) == 1) {//特别地，跳过（已经为1，更新没意义）
				dfs2(now, y);
				continue;
			}
			double pp = g[now] / (1 - edge[i] * (1 - f[y]));
			g[y] = (1 - (1 - pp) * edge[i]) * f[y];
			dfs2(now, y);
		}
}

int main() {
	n = read();
	for (int i = 1, x, y, w; i < n; i++) {
		x = read(), y = read(), w = read();
		add(x, y, (double)w / 100);
		add(y, x, (double)w / 100);
	}
	for (int i = 1; i <= n; i++)
		p[i] = (double)read() / 100;
	dfs1(0, 1);
	g[1] = f[1];
	dfs2(0, 1);
	printf("%.6lf", ans);
}

手工栈版：

#include <stack>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define getchar() (p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 1 << 21, stdin), p1 == p2) ? EOF : *p1++)

char buf[1 << 23], *p1 = buf, *p2 = buf, obuf[1 << 23], *O = obuf;

inline int read() {
	int res = 0, f = 1;
	char c = getchar();
	while (!isdigit(c)) {
		if (c == '-') f = -1;
		c = getchar();
	}
	while (isdigit(c))
		res = res * 10 + (c ^ 48), c = getchar();
	return res * f;
}

struct node {
	int fa, p;
	double e;
};

std::stack<node> s;
int n, tot;
double ans;
int ver[1000001], next[1000001], head[500001], cur[500001];
double edge[1000001], f[500001], p[500001], g[500001];

void add(int u, int v, double w) {
	ver[++tot] = v;
	next[tot] = head[u];
	edge[tot] = w;
	head[u] = tot;
}

void dfs1() {
	s.push((node){0, 1, 0});
	while (s.size()) {
		int p = s.top().p, fa = s.top().fa;
		double e = s.top().e;
		int flag = 0;
		for (int &i = cur[p], y; i; i = next[i])
			if (fa != (y = ver[i])) {
				flag = 1;
				s.push((node){p, y, edge[i]});
				i = next[i];
				break;
			}
		if (!flag) {
			f[fa] *= 1 - e * (1 - f[p]);
			s.pop();			
		}
	}
}

void dfs2() {
	s.push((node){0, 1, 0});
	while (s.size()) {
		int flag = 0;
		int p = s.top().p, fa = s.top().fa;
		for (int &i = cur[p], y; i; i = next[i])
			if (fa != (y = ver[i])) {
				flag = 1;
				if (edge[i] * (1 - f[y]) == 1) {
					s.push((node){p, y, edge[i]});
					i = next[i];
					break;
				}
				double pp = g[p] / (1 - edge[i] * (1 - f[y]));
				g[y] = (1 - (1 - pp) * edge[i]) * f[y];					
				s.push((node){p, y, edge[i]});
				i = next[i];
				break;
			}
		if (!flag) {
			ans += 1 - g[p];
			s.pop();
		}
	}
}

int main() {
	freopen("charger.in", "r", stdin);
	freopen("charger.out", "w", stdout);
	n = read();
	for (int i = 1, x, y, w; i < n; i++) {
		x = read(), y = read(), w = read();
		add(x, y, (double)w / 100);
		add(y, x, (double)w / 100);
	}
	for (int i = 1; i <= n; i++)
		p[i] = (double)read() / 100, f[i] = 1 - p[i];
	memcpy(cur, head, sizeof(head));
	dfs1();
	g[1] = f[1];
	memcpy(cur, head, sizeof(head));
	dfs2();
	printf("%.6lf", ans);
}