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
*/