Code

#include<bits/stdc++.h>

#define ll long long
#define pi pair<int,int>
#define vi vector<int>
#define cpy(x,y,s) memcpy(x,y,sizeof(x[0])*(s))
#define mset(x,v,s) memset(x,v,sizeof(x[0])*(s))
#define all(x) begin(x),end(x)
#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define ary array
#define eb emplace_back
#define IL inline
#define For(i,j,k) for(int i=(j);i<=(k);i++)
#define Fol(i,k,j) for(int i=(k);i>=(j);i--)

using namespace std;

#define B 1400
#define N 500005
#define K 33005
#define mod 998244353

int read(){
    int x=0,f=1;char ch=getchar();
    while(ch<'0' || ch>'9')f=(ch=='-'?-1:f),ch=getchar();
    while(ch>='0' && ch<='9')x=(x<<1)+(x<<3)+(ch^48),ch=getchar();
    return x*f;
}
void write(int x){
    if(x<0)x=-x,putchar('-');
    if(x/10)write(x/10);
    putchar(x%10+'0');
}

void debug(auto &&...x){
    ((cerr<<x<<' '),...);
    cerr<<'\n';
}

IL int pls(int x,int y){return (x+y>=mod?x+y-mod:x+y);}
IL int sub(int x,int y){return (x-y<0?x-y+mod:x-y);}
IL void Add(int &x,int y){x=pls(x,y);}
IL void Dec(int &x,int y){x=sub(x,y);}
IL int mul(int x,int y){return x*1ll*y%mod;}
IL int qp(int x,int y=mod-2){int ans=1;while(y){if(y&1)ans=mul(ans,x);x=mul(x,x);y>>=1;}return ans;}

int id[N],w[N],c[N],a[N];
basic_string<int> e[N];
int vis[B][16],f[B][K],_f[B][K],coef[B][K],ans[N],cnt,tmp[K];
vi S[N];
int n,k;

void FWT(int *g,int opt){
    for(int len=2,mid=1;len<=(1<<k);len<<=1,mid<<=1){
        for(int j=0;j<(1<<k);j+=len){
            for(int k=j;k<j+mid;k++){
                if(!opt)Add(g[k+mid],g[k]);
                else Dec(g[k+mid],g[k]);
            }
        }
    }
}

int calc(int u,int v){
    int res=0;
    if(id[u] && id[v]){
        For(i,0,(1<<k)-1)tmp[i]=mul(_f[id[u]][i],_f[id[v]][i]);
        FWT(tmp,1);
        For(i,0,(1<<k)-1)Add(res,mul(tmp[i],w[i]));
    }
    if(id[v])for(auto val:S[u])Add(res,mul(coef[id[v]][val],w[val]));
    if(id[u])for(auto val:S[v])Add(res,mul(coef[id[u]][val],w[val]));
    for(auto val:S[u])for(auto val2:S[v])Add(res,w[val|val2]);
    return res;
}

void rebuild(int u,int v){
    if(!id[u])id[u]=++cnt;
    for(auto val:S[u])f[id[u]][val]++;S[u].clear();
    if(id[v])For(i,0,(1<<k)-1)f[id[u]][i]+=f[id[v]][i];
    For(i,0,(1<<k)-1)_f[id[u]][i]=f[id[u]][i],coef[id[u]][i]=f[id[u]][i];
    FWT(_f[id[u]],0);
    for(int len=2,mid=1;len<=(1<<k);len<<=1,mid<<=1){
        int z=__lg(len);
        for(int j=0;j<(1<<k);j+=len){
            for(int k=j;k<j+mid;k++){
                int x=coef[id[u]][k],y=coef[id[u]][k+mid];
                coef[id[u]][k]=pls(x,mul(y,c[z]));
                coef[id[u]][k+mid]=pls(x,y);
            }
        }
    }
    For(i,1,k)vis[id[u]][i]=0;
}

void opr(int u,int t){
    if(vis[id[u]][t])return;vis[id[u]][t]=1;
    For(S,0,(1<<k)-1)if(!(S&(1<<(t-1))))Add(f[id[u]][(S|(1<<(t-1)))],f[id[u]][S]),f[id[u]][S]=0;
    For(i,0,(1<<k)-1)_f[id[u]][i]=f[id[u]][i],coef[id[u]][i]=f[id[u]][i];
    FWT(_f[id[u]],0);
    for(int len=2,mid=1;len<=(1<<k);len<<=1,mid<<=1){
        int z=__lg(len);
        for(int j=0;j<(1<<k);j+=len){
            for(int k=j;k<j+mid;k++){
                int x=coef[id[u]][k],y=coef[id[u]][k+mid];
                coef[id[u]][k]=pls(x,mul(y,c[z]));
                coef[id[u]][k+mid]=pls(x,y);
            }
        }
    }
}


void dfs(int x,int fa){
    S[x].pb((1<<(a[x]-1)));ans[x]=c[a[x]];
    for(auto v:e[x]){
        if(v==fa)continue;
        dfs(v,x);
        //operate
        if(id[v])opr(v,a[x]);
        for(auto &val:S[v])val|=(1<<(a[x]-1));
        //calculate
        Add(ans[x],calc(x,v));Add(ans[x],ans[v]);
        //merge
        for(auto val:S[v])S[x].pb(val);vector<int>().swap(S[v]);
        if(id[v])swap(id[x],id[v]);
        if(id[v] || S[x].size()>B)rebuild(x,v);
    }
}

int main(){
    #ifdef EAST_CLOUD
    freopen("a.in","r",stdin);
    freopen("a.out","w",stdout);
    #endif

    n=read(),k=read();
    For(i,1,k)c[i]=read();
    For(i,1,n)a[i]=read();
    For(i,1,n-1){
        int u=read(),v=read();
        e[u]+=v;e[v]+=u;
    }
    For(S,0,(1<<k)-1){
        w[S]=1;
        For(j,1,k)if(S&(1<<(j-1)))w[S]=mul(w[S],c[j]);
    }
    dfs(1,0);
    For(i,1,n)write(ans[i]),putchar(' ');
    return 0;
}

OU 瑞平：你这玩意跟我那个做法感觉没啥区别啊，唯一的区别可能是，常数变大了

写这篇题解的人真是太不牛了。

考虑对原树进行 Top Cluster 分块，这可以将原树的边集不重不漏划分为 $\mathcal O(\sqrt{n})$ 个大小为 $\mathcal O(\sqrt{n})$ 的块。我们先求出 LCA 为 $u$ 的所有答案，然后做树上前缀和即可。

对于单点 $u$，考虑引起贡献的两个点。如果 $u$ 不是这棵簇的界路径上的点，那么引起贡献的 $x, y$ 均与 $u$ 在同一棵簇内，可以直接 $\mathcal O(\left(\sqrt{n}\right)^3)$ 复杂度内暴力解决。

否则，下面的所有点 $u$ 都在某棵簇的界路径上。我们考虑这样的问题：对于当前节点 $u$，有一个数组 $f$，其中 $f_S$ 表示满足 $v \in \textrm{Subtree}(u)$，且 $u \sim v$ 路径上的所有出现过点权集合为 $S$ 的种类数。如果给出 $f$，且仅统计所有 $v \in \textrm{Subtree}(u)$ 的贡献，那么答案是

我们可以换一个角度理解这个事情，假设我们在做 FWT，那么我们在 $f$ 每一维做一个神秘的线性变换到 $g$：$g_0 \gets f_0 + a_uf_1, g_1 \gets a_u(f_0 + f_1)$。答案是 $f_0$。这同样可以 $\mathcal O(k \cdot 2^k)$ 计算答案。

我们假设这棵簇真的就长得像一条链，那么每次 $f$ 向上的时候只需要把所有 $c_x \not\in S$ 的 $f$ 值加到 $f_{S \cup \{c_x\}}$ 里，并且对原 $f_S$ 清零。即可。这条链的长度是 $\mathcal O(\sqrt{n})$ 的，但是有意义的合并只会发生 $k$ 次，而 $2^k + 2^{k-1} + \dots + 1 = \mathcal O(2^k)$，所以复杂度依旧是 $\mathcal O(k \cdot 2^k)$。

现在这棵簇长得不像一条链了！但是在 $u$，我们可以枚举所有簇内的 $x$（其实是包括 $u$ 的）满足 $x$ 的最近簇路径节点为 $u$，这时我们再枚举簇外节点 $y$ 即可补全贡献。所有 $y$ 的信息被携带在数组 $h$ 里，我们只要让 $T$ 表示 $u$ 到 $x$ 路径上所有出现颜色集合，那么 $h_T$ 事实上就是所有可能 $y$ 和 $x$ 的路径答案之和。只要计算出 $h$，贡献是容易简单统计的。

最后，答案唯一可能出现问题的是所有的簇界点，因为它可能同时是若干个簇的上界点，我们在 FWT 暴力合并不同簇的 $f$ 数组时加入贡献即可。

时间复杂度即为 $\mathcal O(n\sqrt{n} + k2^k\sqrt{n})$。