題目傳送門

題解：

　　先註意到一定存在k種顏色，切成k個塊， 然後要求每個塊內的顏色都一樣，所以可以發現同一種顏色一定在同一個塊內，故任意2個相同顏色的最短路勁上的點的顏色都是該顏色。

　　我們可以先把任意相同顏色點對的路徑上的點的顏色都染成這個顏色。 如果發現存在一個點是已經有顏色的話，那麽答案一定為0。

　　至於怎麽顏色， 我們可以暴力往上跑，然後壓縮路徑，和並查集一樣的道理，路過的點都染色且壓縮。

　　這樣就完成了第一部分的處理，接下來就是樹DP了。

　　定義dp[ u ][ 0 ] 的含義是 以u為根的子樹，且u還未被劃入任意一種顏色的方案數。

　　　　dp[ u ][ 1 ] 的含義是 以u為根的子樹，且u已被劃入任意一種顏色的方案數。

　　那麽對於有顏色的點來說，

　　　　dp[u][0] = 0;

對於沒有顏色的點來說：

代碼：

/*
code by: zstu wxk
time: 2019/03/02
Tags: 樹DP
Problem Link: http://codeforces.com/contest/1118/problem/F2
Solve: https://www.cnblogs.com/MingSD/p/10462228.html
*/
#include<bits/stdc++.h>
using namespace std;
 
#define Fopen freopen("_in.txt","r",stdin); freopen("_out.txt","w",stdout);
#define LL long long
#define ULL unsigned LL
#define fi first
#define se second
#define pb push_back
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define lch(x) tr[x].son[0]
#define rch(x) tr[x].son[1]
#define
 
 max3(a,b,c) max(a,max(b,c))
#define min3(a,b,c) min(a,min(b,c))
typedef pair<int,int> pll;
const int inf = 0x3f3f3f3f;
const int _inf = 0xc0c0c0c0;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const LL _INF = 0xc0c0c0c0c0c0c0c0;
const LL mod = 998244353;
const int N = 3e5 + 100;
int n, k, deep[N], fa[N], col[N];
vector<int> vc[N], e1[N];
void dfs(int o, int u){
    deep[u] = deep[o] + 1;
    fa[u] = o;
    for(int v : e1[u]){
        if(v == o) continue;
        dfs(u, v);
    }
}
bool Link(int u, int v, int k){
    int tu = u, tv = v;
    while(u != v){
        if(deep[u] > deep[v]){
            u = fa[u];
            if(col[u] && col[u] != k) return false;
            col[u] = k;
        }
        else {
            v = fa[v];
            if(col[v] && col[v] != k) return false;
            col[v] = k;
        }
    }
    while(tu != u){
        int tt = fa[tu];
        fa[tu] = u;
        tu = tt;
    }
    while(tv != v){
        int tt = fa[tv];
        fa[tv] = v;
        tv = tt;
    }
    return true;
}
LL dp[N][2];
LL t1[N], t2[N], t3[N];
void DFS(int o, int u){
    for(int v : e1[u]){
        if(v == o) continue;
        DFS(u, v);
    }
    if(col[u]) {
        dp[u][0] = 0;
        dp[u][1] = 1;
        for(int v : e1[u]){
            if(v == o) continue;
            dp[u][1] = (dp[u][1] * (dp[v][0] + dp[v][1])) % mod;
        }
    }
    else {
        dp[u][0] = 1;
        dp[u][1] = 0;
        int k = e1[u].size();
        if(k){
            for(int i = 0; i < k; ++i){
                if(e1[u][i] == o)
                    t1[i+1] = 1;
                else {
                    t1[i+1] = dp[e1[u][i]][0] + dp[e1[u][i]][1];
                    dp[u][0] = dp[u][0] * t1[i+1] % mod;
                }
            }
            t2[0] = t3[k+1] = 1;
            for(int i = 1; i <= k; ++i)
                t2[i] = t1[i] * t2[i-1] % mod;
            for(int i = k; i >= 0; --i)
                t3[i] = t1[i] * t3[i+1] % mod;
            for(int i = 0; i < k; ++i){
                int v = e1[u][i];
                if(v == o) continue;
                dp[u][1] = (dp[u][1] + dp[v][1] * t2[i] % mod * t3[i+2] % mod) % mod;
            }
        }
    }
}
void Ac(){
    for(int i = 1, v; i <= n; ++i){
        scanf("%d", &v);
        if(v) vc[v].pb(i);
        col[i] = v;
    }
    for(int i = 1, u, v; i < n; ++i){
        scanf("%d%d", &u, &v);
        e1[u].pb(v);
        e1[v].pb(u);
    }
    dfs(0, 1);
    for(int i = 1; i <= n; ++i){
        for(int j = 1; j < vc[i].size(); ++j){
            if(!Link(vc[i][0],vc[i][j],i)) {
                puts("0");
                return ;
            }
        }
    }
    DFS(0, 1);
    printf("%I64d\n", dp[1][1]);
}
int main(){
    while(~scanf("%d%d", &n, &k)){
        Ac();
    }
    return 0;
}

CodeForces - 1118 F2 Tree Cutting