ABC163F path pass i

Description

给定一棵 $n$ 个点的树，树上每个点有一个颜色。对于每一种颜色，求有多少无向路径至少经过一次该颜色。
$n\leq 10^5$

Solution

正难则反。考虑对于颜色 c ，假设将所有颜色为 c 的点删去，剩下的连通块有 $k$ 个，大小分别为 $s_1,\cdots,s_k$ ，那么颜色 c 的答案为

$\frac{n(n+1)}{2}-\sum_{i=1}^{k}\frac{s_i(s_i+1)}{2}$

得到了一个 $\mathcal{O}(n^2)$ 的做法：暴力枚举颜色再 DFS 。

如何优化？可以发现，$\sum k$ 是 $\mathcal{O}(n)$ 级别的（断掉一个点至多贡献度数个连通块，每个点至多断一次）。对于每一个断后的连通块，其一定有且仅有一个深度最小的节点。我们用这个点来代表这个连通块。除了根以外，所有的点代表且仅代表一个连通块。

这个 trick 很实用。

发现了上述性质，就可以一遍 DFS 解决问题。DFS 同时对于每一种颜色维护一个单调栈，当递归完 $u$ 的子树后，与 $u$ 颜色相同且距离 $u$ 最近的祖先节点 $v$ 其代表的连通块就需要减去 $\text{size}_u$ 。若不存在 $v$ 点，那么这个贡献需要算到根上。即对根单独开一个桶，下标表示颜色，统计所有非根的颜色而计入根的贡献的连通块。还需要注意的是，由于需要删去 $u$ 点，所以我们每遍历一棵 $u$ 的子树，就意味着一个连通块形成了，直接计算即可。

Code

void DFS(int p,int from){
    stk[++top]=p;s[c[p]].push(top);sz[p]=1;
    for(int i=head[p];i;i=nt[i]){
        int v=to[i];
        if(v==from) continue;
        DFS(v,p);sz[p]+=sz[v];
        csz[p]+=sz[v];
        ans[c[p]]+=(1ll*csz[p]*(csz[p]+1)>>1);
        csz[p]=0;
    }
    --top;s[c[p]].pop();
    if(s[c[p]].size()) csz[stk[s[c[p]].top()]]-=sz[p];
    else rsz[c[p]]-=sz[p];
}