21.6.29 t2

tag:树形dp


直接dp,\(O(n^3)\)

唯一的细节可能就是优化空间,最小的不超过 \(\frac n6\),次小的不超过 \(\frac n4\),最大的不超过 \(\frac n2\)


#include<bits/stdc++.h>
using namespace std;

template<typename T>
inline void Read(T &n){
	char ch; bool flag=false;
	while(!isdigit(ch=getchar()))if(ch=='-')flag=true;
	for(n=ch^48;isdigit(ch=getchar());n=(n<<1)+(n<<3)+(ch^48));
	if(flag)n=-n;
}

enum{
	MAXN = 205,
	MOD = 998244353
};

inline int inc(int a, int b){
	a += b;
	if(a>=MOD) a -= MOD;
	return a;
}

inline void iinc(int &a, int b){a = inc(a,b);}
inline void upd(int &a, long long b){a = (a+b)%MOD;}

int f[MAXN/6+5][MAXN/4+5][MAXN/2+5][MAXN], g[MAXN/6+5][MAXN/4+5][MAXN/2+5];
int n, mx[3], id[3], A[MAXN], ps[3];

vector<int>to[MAXN];

inline bool cmp(const int &u, const int &v){return mx[u]<mx[v];}

int sz[3][MAXN], ans;
void dp(int x, int y){
	f[A[x]==0][A[x]==1][A[x]==2][x] = 1;
	sz[A[x]][x] = (mx[id[A[x]]]?1:0);
	for(int v:to[x]) if(v!=y){
		dp(v,x);
		for(int a=0; a<=sz[0][x]; a++)
			for(int b=0; b<=sz[1][x]; b++)
				for(int c=0; c<=sz[2][x]; c++)
					g[a][b][c] = f[a][b][c][x], 
					f[a][b][c][x] = 0;
		for(int a=0; a<=sz[0][v]; a++) for(int aa=0; aa<=sz[0][x] and a+aa<=mx[id[0]]; aa++)
			for(int b=0; b<=sz[1][v]; b++) for(int bb=0; bb<=sz[1][x] and b+bb<=mx[id[1]]; bb++)
				for(int c=0; c<=sz[2][v]; c++) for(int cc=0; cc<=sz[2][x] and c+cc<=mx[id[2]]; cc++)
					upd(f[a+aa][b+bb][c+cc][x],1ll*f[a][b][c][v]*g[aa][bb][cc]);
		for(int i=0; i<3; i++) sz[i][x] = min(sz[i][x]+sz[i][v],mx[id[i]]);
	}
	for(int a=0; a<=sz[0][x]; a++)
		for(int b=0; b<=sz[1][x]; b++)
			for(int c=0; c<=sz[2][x]; c++)
				iinc(ans,f[a][b][c][x]);
	f[0][0][0][x] = 1;
}

int main(){
//	freopen("1.in","r",stdin);
//	freopen("1.out","w",stdout);
	Read(n); for(int i=0; i<3; i++) Read(mx[i]), id[i] = i;
	sort(id,id+3,cmp); for(int i=0; i<3; i++) ps[id[i]] = i;
	for(int i=1; i<=n; i++) Read(A[i]), A[i] = ps[A[i]];
	for(int i=1; i<n; i++){
		int f, t;
		Read(f); Read(t);
		to[f].push_back(t);
		to[t].push_back(f);
	}
	dp(1,0);
	cout<<ans<<'\n';
	return 0;
}

/*
5 1 1 1
0 1 0 0 2
1 2
1 3
3 4
4 5
*/
posted @ 2021-06-29 16:15  oisdoaiu  阅读(20)  评论(0编辑  收藏  举报