这是本文档旧的修订版!
给定一个点权树,点 $i$ 点权为 $a_i$。
树上有向路径 $s\to t$ 被认为是好的当且仅当 $a_s+a_{v_1}x^1+a_{v_2}x^2+\cdots a_tx^k\equiv y(\bmod p)$,其中 $v_1,v_2\cdots$ 为 $s\to t$ 依次经过的点。
三元组 $(u,v,w)$ 被认为是好的当且仅当 $u\to v,u\to w,v\to w$ 三条路径都是好的或都是坏的。(不要求 $u,v,w$ 互异)
问有多少个好的三元组。
构建图 $G$,如果树上有向路径 $u\to v$ 是好的,则 $u\to v$ 连一条黑边,否则 $u\to v$ 连一条白边。
于是图 $G$ 是完全图,所以坏的三元组个数为异色角个数除以 $2$。
假设点 $i$ 黑边入度为 $\text{in}_1$,黑边出度为 $\text{out}_1$,白边入度为 $\text{in}_0$,白边出度为 $\text{out}_0$。
考虑点 $i$ 在三元组 $(u,v,w)$ 中不同位置的异色角贡献。
当 $i$ 作为点 $u$ 时,假设 $u\to v_1$ 是白边,$u\to v_2$ 是黑边,$(u,v_1,v_2),(u,v_2,v_1)$ 是不同的坏的三元组。
于是 $i$ 的异色角贡献为 $2\times \text{out}_0\text{out}_1$。类似的,当 $i$ 作为点 $w$ 时,$i$ 的异色角贡献为 $2\times \text{in}_0\text{in}_1$。
当 $i$ 作为点 $v$ 时,易知 $i$ 的贡献为 $\text{out}_0\text{in}_1+\text{in}_0\text{out}_1$。
于是点 $i$ 的异色角总贡献为 $2\times \text{out}_0\text{out}_1+2\times \text{in}_0\text{in}_1+\text{out}_0\text{in}_1+\text{in}_0\text{out}_1$。
接下来问题转化为怎么计算 $\text{in}_1,\text{out}_1,\text{in}_0,\text{out}_0$,不妨只计算 $\text{in}_1,\text{out}_1$,然后有 $\text{in}_0=n-\text{in}_1,\text{out}_0=n-\text{out}_1$。
考虑点分治,设当前重心为 $\text{rt}$,于是对点对 $(s,t)$,需要判定 $a_s+a_{v_1}x^1+a_{v_2}x^2+\cdots a_tx^k\equiv y(\bmod p)$。
移项,得 $a_{v_k'+1}x^{k'+1}+\cdots +a_tx^k\equiv y-(a_s+a_{v_1}x^1+a_{v_2}x^2+\cdots a_{\text{rt}}x^{k'})(\bmod p)$
设 $\text{pre}(s)=a_s+a_{v_1}x^1+a_{v_2}x^2+\cdots a_{\text{rt}}x^k,\text{suf}(s)=a_{rt}+a_{v_1}x^1+\cdots +a_sx^k$,于是上式化简为
$$ \text{suf}(s)-a_{rt}\equiv (y-\text{pre}(s))x^{-k'}(\bmod p) $$
于是 $\text{dfs}$ 过程中维护 $\text{pre}(s),\text{suf}(s)$ 同时记录 $\text{suf}(s)-a_{rt}$ 和 $(y-\text{pre}(s))x^{-k'}$,最后双指针统计答案。
另外关于两个结点都在同一个子树的情况,直接容斥即可。总时间复杂度为 $O(n\log^2 n)$。