P11648 【MX-X8-T7】「TAOI-3」2236 A.D.

原题链接：。 --- 公元 2236 年 4 月 1 日，收到了 99 封邮件。 这 99 封内容完全相同的邮件的发件人是......她。

Masuko 有一棵 $n$ 个结点的树，以 $1$ 为根，每个结点上都有一个 $[1,k]$ 之间的颜色 $c_i$（$\bm{k \le 15}$），同时给定权值数组 $a_1,a_2,\dots,a_{k}$。 定义两个点 $u,v$ 之间的权值 $f(u,v)$ 如下： - 设 $u=p_1,p_2,\dots,p_m=v$ 是 $u$ 到 $v$ 的最短路径，$f(u,v)=\prod_{i\in\{x|\exists j\in[1,m],c_{p_j}=x\}}a_i$。 即 $u,v$ 最短路径上所有出现过的颜色的权值乘积。 你需要对每个 $u=1,2,3,\dots,n$。求出 $u$ 子树内，所有无序点对的权值和。具体的，假设 $S_u$ 表示所有到根的最短路径上经过 $u$ 的结点组成的集合，你需要输出： $$ \sum_{x,y\in S_u,x\leq y} f(x,y) $$ 答案对 $998244353$ 取模。

第一行，两个正整数 $n,k$。 第二行，$k$ 个非负整数 $a_1,a_2,\dots,a_k$。 第三行，$n$ 个正整数 $c_1,c_2,\dots,c_n$。 接下来 $(n-1)$ 行，每行两个正整数 $u_i, v_i$，表示一条连接点 $u_i$、$v_i$ 的边。

仅一行，$n$ 个非负整数，其中第 $i$ 个表示考虑 $i$ 子树内所有无序点对的权值之和，对 $998244353$ 取模后的结果。

**【样例解释 #1】** - 当 $u=2$ 时，$u$ 子树内只有 $\{2\}$，$f(2,2)=a_{c_2}=1$。 - 当 $u=4$ 时，$u$ 子树内有 $\{2,4\}$，答案为 $f(2,2)+f(4,4)+f(2,4)=1+1+1=3$。 **【数据范围】** **本题采用捆绑测试**。 - 子任务 1（5 分）：$n\leq 1000$； - 子任务 2（10 分）：$n\leq 5000$； - 子任务 3（10 分）：$k=2$； - 子任务 4（10 分）：$v_i=u_i+1$； - 子任务 5（10 分）：$u_i=1$； - 子任务 6（20 分）：$u_i$ 在 $[1,i]$ 中随机生成，$v_i=i+1$； - 子任务 7（15 分）：$n\leq 10^5$； - 子任务 8（10 分）：$n\leq 2\times 10^5$； - 子任务 9（10 分）：无特殊限制。 对于所有数据，保证 $1\leq n\leq 5\times 10^5$，$1\leq k\leq 15$，$0\leq a_i