【PKUWC2018】Minimax

Description

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

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

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

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

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

\[\sum_{i=1}^m i\cdot V_i \cdot D_i ^2 \]

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

Input

第一行一个正整数 \(n\)；

第二行 \(n\) 个整数，第 \(i\) 个整数表示第 \(i\) 个结点的父亲的编号，其中第 \(1\) 个结点的父亲为 \(0\)；

第三行 \(n\) 个整数，若第 \(i\) 个结点没有子结点，则第 \(i\) 个数为它的权值，否则第 \(i\) 个数为 \(p_i\cdot 10000\)，保证 \(p_i\cdot 10000\) 是个正整数。

Output

输出答案。

Solution

据说是去年的签到题啊。话说我最近怎么天天只做签到题

首先考虑DP

先对叶子节点的权值离散化,m为离散化后权值个数

令\(F_i\)表示左子树取到\(i\)这个值得概率, \(G_i\)表示右子树取到\(i\)这个值得概率,\(T_i\)表示根取到\(i\)这个值的概率，易得

\[T_i=F_i * (\sum_{j=1}^{i-1} G_j*p + \sum_{j=i+1}^{m} G_j*(1-p)) + G_i * (\sum_{j=1}^{i-1} F_j*p + \sum_{j=i+1}^{m} F_j*(1-p)) \]

发现这样dp非常慢。。。大概是\(O(n^2)\)，考虑优化。

惊奇的发现我们可以用线段树合并来优化这个dp。。因为每次只和两个数组有关，而且乘上的权值要么是一段前缀和，要么是一段后缀和

所以复杂度变成\(O(n \log n)\)

Code

#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
const int Mod=998244353;
int n,m,b[500000],root[500000],son[500000][2],w[500000],lson[10000010],rson[10000010],exp[10000010],tag[10000010],cnt;
int fpow(int a,int k)
{
	int ans=1;
	while (k)
	{
		if (k&1) ans=1LL*ans*a%Mod;
		a=1LL*a*a%Mod;
		k>>=1;
	}
	return ans;
}
void ins(int x,int t)
{
	exp[x]=1LL*exp[x]*t%Mod;
	tag[x]=1LL*tag[x]*t%Mod;
}
void down(int x)
{
	if (tag[x]!=1)
	{
		ins(lson[x],tag[x]);
		ins(rson[x],tag[x]);
		tag[x]=1;
	}
}
void insert(int &rt,int l,int r,int x)
{
	rt=++cnt;tag[rt]=exp[rt]=1;
	if (l==r) return;
	int mid=(l+r)>>1;
	if (x<=mid) insert(lson[rt],l,mid,x);
	else insert(rson[rt],mid+1,r,x);
}
int Union(int x,int y,int Exp_l,int Exp_r,int Exp)//Exp_l表示取左子树需要乘上的概率,Exp_r表示右子树需要乘上的概率
{
	if (!x)
	{
		ins(y,Exp_r);
		return y;
	}
	if (!y)
	{
		ins(x,Exp_l);
		return x;
	}	
	down(x);down(y);
	int expl_x=exp[lson[x]],expl_y=exp[lson[y]],expr_x=exp[rson[x]],expr_y=exp[rson[y]];
	lson[x]=Union(lson[x],lson[y],(Exp_l+1LL*expr_y*(1-Exp+Mod)%Mod)%Mod,(Exp_r+1LL*expr_x*(1-Exp+Mod)%Mod)%Mod,Exp);
	rson[x]=Union(rson[x],rson[y],(Exp_l+1LL*expl_y*Exp%Mod)%Mod,(Exp_r+1LL*expl_x*Exp%Mod)%Mod,Exp);
	exp[x]=(exp[lson[x]]+exp[rson[x]])%Mod;
	return x;
}
void solve(int x)
{
	if (!son[x][0])
	{
		insert(root[x],1,m,lower_bound(b+1,b+m+1,w[x])-b);
		return;
	}
	solve(son[x][0]);
	if (!son[x][1])
	{
		root[x]=root[son[x][0]];
		return;
	}
	solve(son[x][1]);
	root[x]=Union(root[son[x][0]],root[son[x][1]],0,0,w[x]);
} 
int calc(int rt,int l,int r)
{
	if (l==r)  return 1LL*l*b[l]%Mod*exp[rt]%Mod*exp[rt]%Mod;
	int mid=(l+r)>>1;
	down(rt);
	return (calc(lson[rt],l,mid)+calc(rson[rt],mid+1,r))%Mod;
}
int main()
{
	scanf("%d",&n);
	for (int i=1;i<=n;i++)
	{
		int x;
		scanf("%d",&x);
		son[x][son[x][0]?1:0]=i;
	}
	for (int i=1;i<=n;i++)
	{
		int x;
		scanf("%d",&x);
		if (son[i][0]) w[i]=1LL*x*fpow(10000,Mod-2)%Mod;
		else w[i]=x,b[++m]=x;
	}
	sort(b+1,b+m+1);//离散化
	solve(1);
	printf("%d\n",calc(root[1],1,m));
	return 0;
}