树上连通块dp速记

· · 个人记录

不知道为什么突然就想写点什么，先放两个题目
P2664 树上游戏
[ABC163F] path pass i
也是偶然做到这两个题目，发现有一些相同点，于是就总结一下。

首先，这两个题目都和颜色数有关。
于是想到第一个trick，将每种不同的颜色分开考虑。并从反面的角度（容斥）去考虑

有了这第一个trick以后，这两题都转化为一些关于数上连通块问题，我抽象出来一个简单的模型：

给定一颗涂了颜色的树。求去除某种颜色的节点和相连的边后，每个连通块的大小。

对于每一个节点 i ，设 i 号节点的颜色为 c_i ，求 i 的子树内根节点颜色为 c_i 的极大子树的信息。这样描述不好理解，画个图理解一下。

图中的根节点为当前的 i ，红色点代表和根节点颜色一样的，白色的为不一样的。其中有两个极大的根的颜色和 i 相同的子树，用蓝色框框标出的。

为什么叶子节点的红色节点不算子树？因为要保证 极大，这个节点已经是另一个根的子树了，所以不是极大。

可能有同学会问，这个东西的定义这么复杂，求它有什么用呢？
设 sz_i 为 i 的子树大小， k_{i,v} 表示以 i 为根，根节点在 v 子树内的极大子树的 sz 和。那么显然，去掉 c_i 这种颜色的点后的连通块信息我们就能计算了。
对于去掉 c 后形成的连通块，如果不包括1号节点，连通块中至少有一个点的父亲为 i c_i=c , 或者包括了1号节点。

我们不妨先计算不包括一号节点的连通块。我们dfs每一个节点，对于这个节点的所有儿子节点，如果它和父亲节点颜色不同，他们一定被包括在一个连通块内，而且这样计算是不重不漏的。 

void dfs(int u, int fa) {
    sz[u]=1;
    ll tmp=cnt[col[u]], c=col[u];
    for(int v:e[u]) {
        if(v==fa) continue ;
        int t=cnt[c];
        dfs(v, u);
        int Ad=cnt[c]-t;
        sz[u]+=sz[v];
        ans[c]-=Sum(sz[v]-Ad);
    }
    cnt[c]=tmp+sz[u];
}

我们记一个数组 cnt_c 表示颜色 c 的极大子树大小。当dfs到当前节点 u ，在dfs子树之前先用零时变量 t 记录当前的 cnt_c 的大小，然后dfs完成以后增加的量 \Delta cnt_c 即为 当前点的颜色的极大子树的大小。最后将 cnt_c 的值更新为 t+sz_u 。
这样做是为什么呢当前 u 节点内的子树不再极大了，因此不用记录，子树外的点不会有影响，当前点的子树是一颗极大的子树，所以令 t+sz_u \rightarrow cnt_c 这样就能在 o(n) 的时间找出所有连通块的信息了。
还有一点，如果连通块包括1号节点，那么我们以下的dfs是找不出来的，但是我们可以增加一个0号节点，他的颜色是所有颜色，连向1号节点。这个0号节点不用建出来，只需要特殊处理一下就行了。 

    for(int i=1; i<=n; i++)
        printf("%lld\n", ans[i]-Sum(n-cnt[i]));

其中 n-cnt_i 就是包括1号节点的连通块。