dsu on the tree 学习笔记

dsu貌似就是启发式合并的意思，大家都知道这个东西是很神奇的，把n^2级别的复杂度降到nlogn，特别厉害的。

考虑如下问题：一棵树，问以每个点为根的子树中，和它颜色相同的有多少个节点（n<=10w,颜色数因为可以离散所以不重要）。

主席树+dfs序是很显然的，不过二维数据结构太高级了。考虑一下别的做法。先从暴力想起吧，O(n^2)暴力大家都会。这里给出一个不是那么暴力的nlogn做法（感觉好像是已知的最优复杂度了），考虑n^2暴力的时候，我们使用两层dfs解决，第一层是遍历原树所有节点，然后在每一个节点进入第二层dfs去遍历他的子数，并且记录一个cnt数组，然后根据cnt数组回答询问（这里只用一个变量记录显然也是可以的，只是这样更容易与下文方法衔接）。现在我们做一个改进，按道理来说，我们每次第二层dfs完某棵子树就应该把cnt清零，但是我们不这样做，我们按从儿子到父亲的顺序进行第二层的dfs，也就是说，如果父亲进入了第二层dfs，那么他的所有儿子一定dfs完了。我们按照轻重链剖分的方法找出重儿子，然后在递归的时候不清除重儿子的cnt标记，这要求我们在对节点u进行第二次dfs之前，先dfs所有轻儿子，然后dfs重儿子，保留cnt数组的信息，同时，每个节点用一个vector来保存子树的所有颜色信息，意即vector的大小即为子树大小，顺序任意，然后u继承它重儿子的vector，暴力将另外的子树的信息插入这个vector，更新cnt信息。这样的复杂度是nlogn，重点在于选重儿子，这样的话，每暴力将一个点合并到重儿子的信息上，它所在集合的大小至少会×2，那么最多合并logn次，有n个节点，于是复杂度为nlogn。

代码如下（代码采用C++11标准，auto那个表示遍历vector中的每个元素：

vector<int> *vec[maxn];//为了方便继承重儿子的vector，这里要用指针，也保证了空间复杂度不超过nlogn，分析类似时间复杂度的分析
int cnt[maxn];
void dfs(int v, int p, bool keep){
    int mx = -1, bigChild = -1;
    for(auto u : g[v])
       if(u != p && sz[u] > mx)
           mx = sz[u], bigChild = u;
    for(auto u : g[v])
       if(u != p && u != bigChild)
           dfs(u, v, 0);
    if(bigChild != -1)//最后来dfs重儿子，保证cnt数组性质
        dfs(bigChild, v, 1), vec[v] = vec[bigChild];
    else
        vec[v] = new vector<int> ();
    vec[v]->push_back(v);
    cnt[ col[v] ]++;
    for(auto u : g[v])
       if(u != p && u != bigChild)
           for(auto x : *vec[u]){
               cnt[ col[x] ]++;
               vec[v] -> push_back(x);
           }
    //now (*cnt)[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
    // note that in this step *vec[v] contains all of the subtree of vertex v.
    if(keep == 0)//不是重儿子的需要暴力还原
        for(auto u : *vec[v])
            cnt[ col[u] ]--;
}

这个方法其实比较麻烦，（对于原文来说，我觉得map那个方法麻烦到天际），稍微改进一下，不用vector记录信息了，如果我需要还原，直接把整棵子树（因为这时候重儿子的信息也在里面）暴力dfs一遍还原，为什么不会T？考虑一下，如果我某一个点被暴力dfs到，那么它一定不在当前这个子树根节点的重儿子内，那么整棵子树的大小至少是他所在那一棵小子树的大小的二倍，显然的是至多只会有logn个节点作为子树根节点把它dfs到。所以时间复杂度是对的。

代码如下：

int cnt[maxn];
bool big[maxn];
void add(int v, int p, int x){
    cnt[ col[v] ] += x;
    for(auto u: g[v])
        if(u != p && !big[u])
            add(u, v, x)
}
void dfs(int v, int p, bool keep){
    int mx = -1, bigChild = -1;
    for(auto u : g[v])
       if(u != p && sz[u] > mx)
          mx = sz[u], bigChild = u;
    for(auto u : g[v])
        if(u != p && u != bigChild)
            dfs(u, v, 0);  // run a dfs on small childs and clear them from cnt
    if(bigChild != -1)
        dfs(bigChild, v, 1), big[bigChild] = 1;  // bigChild marked as big and not cleared from cnt
    add(v, p, 1);
    //now cnt[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
    if(bigChild != -1)
        big[bigChild] = 0;
    if(keep == 0)
        add(v, p, -1);
}

总是要写两层dfs，不太直观，虽然第二层dfs确实很简单，，不过我们再来优化一下代码实现，说来也简单，只要用在dfs序列上遍历一个区间来代替原来的第二层dfs即可。st[u]和ft[u]是dfs序上以u为根的子树的左右端点，这样就好写的多了。

代码：

int cnt[maxn];
void dfs(int v, int p, bool keep){
    int mx = -1, bigChild = -1;
    for(auto u : g[v])
       if(u != p && sz[u] > mx)
          mx = sz[u], bigChild = u;
    for(auto u : g[v])
        if(u != p && u != bigChild)
            dfs(u, v, 0);  // run a dfs on small childs and clear them from cnt
    if(bigChild != -1)
        dfs(bigChild, v, 1);  // bigChild marked as big and not cleared from cnt
    for(auto u : g[v])
    if(u != p && u != bigChild)
        for(int p = st[u]; p < ft[u]; p++)
        cnt[ col[ ver[p] ] ]++;
    cnt[ col[v] ]++;
    //now cnt[c] is the number of vertices in subtree of vertex v that has color c. You can answer the queries easily.
    if(keep == 0)
        for(int p = st[v]; p < ft[v]; p++)
        cnt[ col[ ver[p] ] ]--;
}

参考文献：http://codeforces.com/blog/entry/44351