【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;
}