题解：P11648「TAOI-3」2236 A.D.

\(\text{Link}\)

题意

给你一颗 \(n\) 个结点的树，结点 \(1\) 为根，结点有 \([1,k]\) 的颜色 \(c_i\)。给定序列 \(a_{1\sim k}\)，定义两个结点 \(x,y\) 之间的权值 \(f(x,y)\) 为所有在 \(x\to y\) 的简单路径上出现过的颜色 \(c\) 的 \(a_c\) 之积。

对 \(\forall i\in[1,n]\) 求出 \(i\) 子树内的点两两的权值之和，对 \(998244353\) 取模。

\(n\le 5\times 10^5\)，\(k\le 15\)。

题解

直接对每个颜色集合预处理出其权值 \(v_s\)，考虑计算 \(x\) 的不同子树 \(i,j\) 间的答案。

维护子树 \(i\) 内每个点到 \(i\) 的路径颜色集合 \(S_i\)，需要支持两种操作：将集合 \(S_i\) 内所有元素并上 \(c_x\)；查询 \(S_i,S_j\) 之间两两元素 \(p,q\) 的 \(v_{p\cup q}\) 之和，再合并这两个集合。

一眼可以看出两个暴力，分别是直接枚举和 FWT 卷积，于是考虑根号分治。设置阈值 \(b\)，若 \(S_i,S_j\) 均有 \(\ge b\) 个元素便直接卷积；若 \(S_i,S_j\) 均有 \(\le b\) 个元素则直接枚举；否则对 \(S_i\) 内每个元素 \(p\) 计算 \(\sum_{q\in S_j}cnt_qv_{p\cup q}\)，这可以采用 All Pairs Similarity P 的做法，进行 \(O(2^kk)\) 预处理 \(O(1)\) 回答。

---

具体地，设 \(f_p\) 表示 \(p\) 的答案，集合的包含、并集等条件引出一类做法，考虑设计一辅助状态 \(g_s\) 满足：

\[f_s=\sum_{T⊇S} g_t \]

仅考虑一个元素 \(q\in S_j\)，我们希望构造出 \(g_t\) 满足其高维后缀和为 \(f_s=v_{s\cup q}\)。那么我们对 \(v_{0\sim 2^k-1}\) 做高维后缀差分得到数组 \(w_{0\sim 2^k-1}\) 满足：

\[v_s=\sum_{T⊇S} w_t \]

构造 \(g_t=[q\sube T]w_{t}\) 即可，容易验证，所有同时是 \(p\) 和 \(q\) 的超集的数 \(t\) 都会作出 \(w_{t}\) 的贡献，它们的和正好等于 \(v_{p\cup q}\)。

每次预处理都需要做一次高维前缀和和一次高维后缀和，时间复杂度为 \(O(2^kk)\)，单次查询 \(O(1)\)。

---

注意每次计算完还需要合并，我们对每个集合分别维护一个 \(\le b\) 个元素的零散部分和非零散部分的信息，合并后若零散部分大小 \(\ge b\) 直接将其重构入非零散部分即可。将集合内的数取并可以直接暴力对零散部分做，维护非零散部分的取并标记。

重构与卷积次数分别只有 \(O(n/b)\) 次，零散部分的枚举量为 \(\sum_{i,j}\min(siz_i,b)\min(siz_j,b)\) 即 \(O(nb)\)。

取 \(b=\sqrt{2^kk}\)，时间复杂度 \(O(n\sqrt{2^kk})\)。

---