Crystal's Code

Crystal's Code

#include<bits/stdc++.h>
using namespace std;
#define int long long 
const int maxn=1e3;
const int mod=998244353;
struct Node{
	int a,b,c,d;
	Node(){}
	Node(int ia,int ib,int ic,int id){
		a=ia,b=ib,c=ic,d=id;
	}
	Node operator+(Node y){
		return Node((a+y.a)%mod,(b+y.b)%mod,(c+y.c)%mod,(d+y.d)%mod);
	}
	void operator+=(Node y){
		a=(a+y.a)%mod;
		b=(b+y.b)%mod;
		c=(c+y.c)%mod;
		d=(d+y.d)%mod;
	}
	Node operator*(Node y){
		Node t;
		t.a=(a*y.a)%mod;
		t.b=(a*y.b+b*y.a)%mod;
		t.c=(a*y.c+2*b*y.b+c*y.a)%mod;
		t.d=(a*y.d+3*b*y.c+3*c*y.b+d*y.a)%mod;
		return t;
	}
	inline void Init(int x){
		a=bool(x),b=x,c=b*x%mod,d=c*x%mod;
	}
	inline int Get(int k){
		if(k==1) return b;
		else if(k==2) return c;
		else return d;
	}
};
int n,K;
vector<int> E[maxn+5];
Node f[maxn+5][maxn+5][2];
Node tmp[maxn+5];
int ans;
signed main(){
	int t;
	scanf("%lld%lld",&n,&K);
	for(int i=2;i<=n;i++){
		scanf("%lld",&t);
		E[t].push_back(i);
	}
	for(int u=n;u>=1;u--){
		for(int x=1;x<=n;x++){
			f[u][x][0].Init(x);
			f[u][x][1].Init(0);
			for(int v:E[u]){
				tmp[0].Init(0);
				tmp[1].Init(0);
				for(int y=max(1ll,x-1);y<=min(x+1,n);y++){
					if(y==x+1){
						tmp[0]+=(f[v][y][0]+f[v][y][1])*f[u][x][0];
						tmp[1]+=(f[v][y][0]+f[v][y][1])*f[u][x][1];
					}
					else if(y==x-1){
						tmp[1]+=f[v][y][1]*(f[u][x][1]+f[u][x][0]);
					}
					else{
						tmp[0]+=f[v][y][1]*f[u][x][0];
						tmp[1]+=f[v][y][1]*f[u][x][1];
						if(x==1&&y==1){
							tmp[1]+=(f[v][y][0]+f[v][y][1])
							*(f[u][x][1]+f[u][x][0]);
						}
					}
				}
				f[u][x][0]=tmp[0];
				f[u][x][1]=tmp[1];
			}
		}
	}
	for(int i=1;i<=n;i++) ans=(ans+f[1][i][1].Get(K))%mod;
	ans=2*ans%mod;
	printf("%lld",ans);
	return 0;
}