【JZOJ 省選模擬】6655.三國學者

題目

Description

Input

Output

Sample Input
5 5 3
1 2 3 3
3 2 2 1 2
3 2 2 2 2

Sample Output
12

Data Constraint
對於所有數據，有Ai<L<=N<=10^6, M<=2*10^6, fai<i,vi<L.

思路

先考慮暴力dp
設 f[i][j]爲從根節點到i的路徑上有多少條結尾權值爲j的路徑，最後加上所有wx=1的節點的f[x][2]就行。

考慮優化
我們發現每到⼀個節點，f[i] 相對於 f[f[ai]]只改變了一個位置。那麼我們只需要通過dfs來維護一個全局的dp數組，來表示出每個節點的f即可。
即在進入個點時修改f ，退出時把這個點的影響消除。

現在的問題是如何快速地計算gx
看起來不能用桶直接從子樹搞上來，因爲涉及到合併的問題。
但是大佬又提供了一種神奇的思路：直接將g值加入桶中，從x往下遞歸之前先記一下，遞歸回來之後作差，就可以得到子樹對gx的貢獻。

代碼

#include<bits/stdc++.h>
#define N 1000005
#define ll long long 
#define mod 1000000007
using namespace std;

int n,m,L,lim,fa[N],a[N],b[N],now;
int cnt,e[N],nx[N],ls[N];
ll sum,vf[N],vg[N],Ans,d[N];

void dfs(int x)
{
	if(x<=lim) (vf[b[x]]+=vf[b[x]+1])%=mod; 
	ll g0=vg[a[x]-1],g1=vg[b[x]-1];
	for(int i=ls[x]; i; i=nx[i]) dfs(e[i]);
	g0=vg[a[x]-1]-g0,g1=vg[b[x]-1]-g1;
	if(a[x]==1) g0=1; if(b[x]==1) g1=1;
	if(x<=lim) {
		(vf[b[x]]-=vf[b[x]+1])%=mod;
		d[x]=g1%mod*vf[b[x]+1]-g0%mod*vf[a[x]+1];
	}
	(vg[a[x]]+=g0)%=mod;
}

int main(){
	freopen("scholar.in","r",stdin);
	freopen("scholar.out","w",stdout);
	scanf("%d%d%d",&n,&m,&L);
	for(int i=2; i<=n; i++) scanf("%d",&fa[i]);
	for(int i=n; i>=2; i--) cnt++,e[cnt]=i,nx[cnt]=ls[fa[i]],ls[fa[i]]=cnt;
	for(int i=1; i<=n; i++) scanf("%d",&a[i]);
	
	now=0; int tp=0; vf[L+1]=1;
	while(m)
	{
		lim=min(n,m); m-=lim;
		for(int i=1; i<=lim; i++) scanf("%d",&b[i]);
		if(tp) for(int i=1; i<=n; i++) vg[a[i]]=0; else tp=1;
		dfs(1);
		sum=vg[L];
		for(int i=1; i<=lim; i++)
		{
			(sum+=d[i])%=mod,a[i]=b[i];
			Ans+=sum*(now+i)%mod;
		}
		now+=lim;
	}
	printf("%lld",(Ans%mod+mod)%mod);
}