LOJ #2537. 「PKUWC 2018」Minimax (线段树合并 优化dp)

题意

小 \(C\) 有一棵 \(n\) 个结点的有根树，根是 \(1\) 号结点，且每个结点最多有两个子结点。

定义结点 \(x\) 的权值为：

1.若 \(x\) 没有子结点，那么它的权值会在输入里给出，保证这类点中每个结点的权值互不相同。

2.若 \(x\) 有子结点，那么它的权值有 \(p_x\) 的概率是它的子结点的权值的最大值，有 \(1-p_x\) 的概率是它的子结点的权值的最小值。

现在小 \(C\) 想知道，假设 \(1\) 号结点的权值有 \(m\) 种可能性，权值第 \(i\) 小的可能性的权值是 \(V_i\) ，它的概率为 \(D_i(D_i>0)\) ，求：

\[\displaystyle \sum _{i=1} ^ {m} i \cdot V_i \cdot D_i^2
\]

你需要输出答案对 \(998244353\) 取模的值。

对于 \(40\%\) 的数据，有 \(1\leq n\leq 5000\) ；

对于 \(100\%\) 的数据，有 \(1\leq n\leq 3\times 10^5, 1\leq w_i\leq 10^9\)。

题解

首先考虑 \(O(n^2)\) 的 dp , 令 \(dp_{u,i}\) 为 \(u\) 号点 , 取到排名为 \(i\) 权值的概率 .

这个应该比较容易转移 , 考虑枚举一个儿子取的值 , 然后对于它的贡献 就分为它最小和它最大的两种去计数就行了 .

令 \(ls, rs\) 为 \(u\) 的左/右儿子 , 枚举两个子树的取值 \(i, j\) , 令 \(coef = dp[ls][i] * dp[rs][j]\) 方程就是 :

\[coef * prob[u] \to dp[u][max(i, j)] \\
coef * (1 - prob[u]) \to dp[u][min(i, j)]
\]

如果我们只枚举有效状态那么就是 \(O(n ^ 2)\) 的啦，因为对于一对点对他们的贡献只会在 \(lca\) 合并。

---

然后考虑优化 , 类似于这种状态数与 \(size_u\) 有关的 dp .

常常可以考虑 线段树合并 or 启发式合并 来优化时间复杂度 .

一开始想直接 启发式合并 在线段树上操作 发现细节好多 而且不好维护 ... 然后就弃掉了 看了一波 LOJ 最短代码 qwq

诶 好像很好写啊 , 原来直接线段树合并就行了 . qwq

考虑维护一颗线段树 , 每个点维护两个值 \(sumv, mult\) 代表 区间和 以及 区间乘法的标记 .

然后每个叶子 代表一个 dp 值 , 然后每个区间就可以维护这段区间的 dp 值之和 .

我们一边合并一边算到当前区间 , 对于两个线段树 dp 值存在的贡献 \(sumx, sumy\) (也就是前面方程中需要乘上后面的两个东西) .

如果当前区间只有一个子树 , 打下乘法标记 , 直接返回就行了 . 否则继续递归下去合并解决 .

时间复杂度就是 $ O(\sum_{i=1}^{n} minsize) = O(n \log n)$ .

这是因为每个点合并上去 大小至少翻倍 . 意味着每个点最多被计算 \((\log n)\) 次 , 最后复杂度就是 \(O(n \log n)\) .

代码

\[40pts
\]

#include <bits/stdc++.h>
#define For(i, l, r) for(register int i = (l), i##end = (int)(r); i <= i##end; ++i)
#define Fordown(i, r, l) for(register int i = (r), i##end = (int)(l); i >= i##end; --i)
#define Set(a, v) memset(a, v, sizeof(a))
using namespace std;

inline bool chkmin(int &a, int b) {return b < a ? a = b, 1 : 0;}
inline bool chkmax(int &a, int b) {return b > a ? a = b, 1 : 0;}

inline int read() {
    int x = 0, fh = 1; char ch = getchar();
    for (; !isdigit(ch); ch = getchar() ) if (ch == '-') fh = -1;
    for (; isdigit(ch); ch = getchar() ) x = (x << 1) + (x << 3) + (ch ^ 48);
    return x * fh;
}

void File() {
#ifdef zjp_shadow
	freopen ("2537.in", "r", stdin);
	freopen ("2537.out", "w", stdout);
#endif
}

typedef long long ll;
const ll Mod = 998244353;
ll fpm(ll x, int power) {
	ll res = 1;
	for (; power; power >>= 1, (x *= x) %= Mod)
		if (power & 1) (res *= x) %= Mod;
	return res;
}

typedef long long ll;
const int N = 5100, inv = fpm(10000, Mod - 2);
int n, fa[N], ch[N][2], tot[N], val[N], rk[N], Leaf;
ll dp[N][N], p[N], Pre[N][2], Suf[N][2];

#define ls(o) ch[o][0]
#define rs(o) ch[o][1]
void Dp(int u) {
	if (!u) return ; Dp(ls(u)); Dp(rs(u));
	if (!tot[u]) { dp[u][rk[u]] = 1; }
	else if (tot[u] == 1) {
		For (i, 1, Leaf) dp[u][i] = dp[ls(u)][i];
	} else {
		For (son, 0, 1) {
			For (i, 1, Leaf)
				Pre[i][son] = (Pre[i - 1][son] + dp[ch[u][son]][i]) % Mod;
			Fordown (i, Leaf, 1)
				Suf[i][son] = (Suf[i + 1][son] + dp[ch[u][son]][i]) % Mod;
		}
		For (i, 1, Leaf) For (son, 0, 1) {
			(dp[u][i] += dp[ch[u][son]][i] * Pre[i - 1][son ^ 1] % Mod * p[u] % Mod
			 + dp[ch[u][son]][i] * Suf[i + 1][son ^ 1] % Mod * (Mod + 1 - p[u]) % Mod) %= Mod;
		}
	}
}

int main () {
	File();
	n = read();
	For (i, 1, n) fa[i] = read(), ch[fa[i]][tot[fa[i]] ++] = i;
	For (i, 1, n)
		if (!tot[i]) rk[i] = val[++ Leaf] = read();
		else p[i] = 1ll * read() * inv % Mod;

	sort(val + 1, val + Leaf + 1);
	For (i, 1, n) rk[i] = lower_bound(val + 1, val + Leaf + 1, rk[i]) - val;

	Dp(1);

	ll ans = 0;
	For (i, 1, n)
		(ans += 1ll * i * val[i] % Mod * dp[1][i] % Mod * dp[1][i] % Mod) %= Mod;
	printf ("%lld\n", ans);
    return 0;
}

\[100pts
\]

#include <bits/stdc++.h>
#define For(i, l, r) for(register int i = (l), i##end = (int)(r); i <= i##end; ++i)
#define Fordown(i, r, l) for(register int i = (r), i##end = (int)(l); i >= i##end; --i)
#define Set(a, v) memset(a, v, sizeof(a))
using namespace std;

inline bool chkmin(int &a, int b) {return b < a ? a = b, 1 : 0;}
inline bool chkmax(int &a, int b) {return b > a ? a = b, 1 : 0;}

inline int read() {
    int x = 0, fh = 1; char ch = getchar();
    for (; !isdigit(ch); ch = getchar() ) if (ch == '-') fh = -1;
    for (; isdigit(ch); ch = getchar() ) x = (x << 1) + (x << 3) + (ch ^ 48);
    return x * fh;
}

void File() {
#ifdef zjp_shadow
	freopen ("2537.in", "r", stdin);
	freopen ("2537.out", "w", stdout);
#endif
}

typedef long long ll;
const ll Mod = 998244353;
ll fpm(ll x, int power) {
	ll res = 1;
	for (; power; power >>= 1, (x *= x) %= Mod)
		if (power & 1) (res *= x) %= Mod;
	return res;
}

typedef long long ll;

const int N = 3e5 + 1e3, inv = fpm(10000, Mod - 2);
int val[N], rk[N];

#define ls(o) ch[o][0]
#define rs(o) ch[o][1]
#define lson ls(o), l, mid
#define rson rs(o), mid + 1, r
const int Maxnode = 6e6 + 1e3;
#define Mult(o, val) (sumv[o] *= (val)) %= Mod, (mult[o] *= (val)) %= Mod;
struct Segment_Tree {
	int rt[Maxnode], ch[Maxnode][2], Size; ll sumv[Maxnode], mult[Maxnode];

	inline void push_up(int o) { sumv[o] = (sumv[ls(o)] + sumv[rs(o)]) % Mod; }

	inline void push_down(int o) {
		if (mult[o] <= 1) return ; Mult(ls(o), mult[o]); Mult(rs(o), mult[o]); mult[o] = 1;
	}

	void Update(int &o, int l, int r, int up, ll uv) {
		if (!o) o = (++ Size); mult[o] = 1;
		if (l == r) { (sumv[o] += uv) %= Mod; return ; } int mid = (l + r) >> 1;
		push_down(o); if (up <= mid) Update(lson, up, uv); else Update(rson, up, uv); push_up(o);
	}

	int Merge(int x, int y, ll sumx, ll sumy, ll probmax, ll probmin) {
		if (!y) { Mult(x, sumy); return x; }
		if (!x) { Mult(y, sumx); return y; }
		push_down(x); push_down(y);
		ll x0 = sumv[ls(x)], x1 = sumv[rs(x)], y0 = sumv[ls(y)], y1 = sumv[rs(y)];
		ls(x) = Merge(ls(x), ls(y), (sumx + probmin * x1) % Mod, (sumy + probmin * y1) % Mod, probmax, probmin);
		rs(x) = Merge(rs(x), rs(y), (sumx + probmax * x0) % Mod, (sumy + probmax * y0) % Mod, probmax, probmin);
		push_up(x); return x;
	}

	inline ll Calc(int o, int l, int r) {
		if (l == r) return 1ll * l * val[l] % Mod * sumv[o] % Mod * sumv[o] % Mod;
		int mid = (l + r) >> 1; push_down(o);
		return (Calc(lson) + Calc(rson)) % Mod;
	}
} T;

int n, fa[N], ch[N][2], tot[N], Leaf;
ll p[N];

void Dp(int u) {
	if (!u) return ; Dp(ls(u)); Dp(rs(u));
	if (!tot[u]) T.Update(T.rt[u], 1, Leaf, rk[u], 1);
	else if (tot[u] == 1) T.rt[u] = T.rt[ls(u)];
	else T.rt[u] = T.Merge(T.rt[ls(u)], T.rt[rs(u)], 0, 0, p[u], (Mod + 1 - p[u]) % Mod);
}

int main () {
	File();
	n = read();
	For (i, 1, n) fa[i] = read(), ch[fa[i]][tot[fa[i]] ++] = i;
	For (i, 1, n)
		if (!tot[i]) rk[i] = val[++ Leaf] = read();
		else p[i] = 1ll * read() * inv % Mod;

	sort(val + 1, val + Leaf + 1);
	For (i, 1, n) rk[i] = lower_bound(val + 1, val + Leaf + 1, rk[i]) - val;

	Dp(1);

	printf ("%lld\n", T.Calc(T.rt[1], 1, Leaf));
    return 0;
}