hihoCoder #1073 光棍节

树上的路径统计问题。树的点分治。

Implementation

#include <bits/stdc++.h>
using namespace std;

const int N=1e5+5;
int a[N];
vector<int> g[N];
int size[N];
bool removed[N];

pair<int,int> centroid(int u, int f, const int &tot){
    size[u]=1;
    pair<int,int> res={INT_MAX, 0};
    int ma=0; // the max size of connected components after deleting u
    for(auto v: g[u]){
        if(v!=f && !removed[v]){
            res=min(res, centroid(v, u, tot));
            size[u]+=size[v];
            ma=max(ma, size[v]);
        }
    }
    ma=max(ma, tot-size[u]);
    return min(res, {ma, u}); // error-prone
}


map<int,long long> res;
unordered_map<int,int> sum, tmp;

void dfs(int u, int f, int gcd){
    res[gcd]++;
    tmp[gcd]++;
    for(auto v: g[u])
        if(v!=f && !removed[v])
            dfs(v, u, __gcd(gcd, a[v]));
}

void DC(int u, int n){
    int root=centroid(u, u, n).second;
    removed[root]=true;
    res[a[root]]++;


    sum.clear();    // error-prone

    for(auto v: g[root])    //error-prone
        if(!removed[v]){
            tmp.clear();
            dfs(v, root, __gcd(a[root], a[v]));
            for(auto x: tmp){
                for(auto y: sum)
                    res[__gcd(x.first, y.first)]+=(long long)x.second*y.second;
                // sum[x.first]+=x.second; // !error
            }
            for(auto x: tmp)
                sum[x.first]+=x.second;
        }

    for(auto v: g[root]){
        if(!removed[v]){
            int tot=size[v]<size[root]?size[v]:n-size[root]; // 这种写法比较繁琐
            DC(v, tot);
        }
    }

}

int main(){
    int n;
    scanf("%d", &n);
    for(int i=1; i<=n; i++)
        scanf("%d", a+i);
    for(int i=1; i<n; i++){
        int u, v;
        scanf("%d%d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }

    DC(1, n);

    for(auto x: res)
        printf("%d %lld
", x.first, x.second);

    return 0;
}

点分治的实现细节

1. 求重心时需要把连通块（子树）的大小作为参数。
2. 【统计包含根节点（重心）的路径】必须在【递归处理子树】之前进行。由于删点标记是在递归之前打上的，【递归处理子树】后，各子树（即删除重心后所剩的个各连通块）就被破坏了。