洛谷2664树上游戏-点分治

link：https://www.luogu.com.cn/problem/P2664
lrb 有一棵树，树的每个节点有个颜色。给一个长度为 \(n\) 的颜色序列，定义 \(s(i,j)\) 为 \(i\) 到 \(j\) 的颜色数量。以及

现在他想让你求出所有的 \(sum_i\)。

一个暴力的想法：因为是求和，所以可以拆开算贡献。枚举每个颜色 \(c\)，将颜色 \(c\) 的点拿出来，会把原树划分成若干个连通块，对每个点的贡献即为 \(n-sz_i\) ，其中 \(sz_i\) 表示 \(i\) 所属连通块大小。连通块用并查集维护，这样是 \(O(n^2)\)的。

考虑点分治，对于当前的分治中心 \(x\)，需要考虑：

• 以 \(x\) 为某个端点，延伸到某个子节点的产生的答案。
• 以 \(x\) 为LCA，从某个子树中的 \(y\) 出发（或直接从 \(x\) 出发）对另一子树中 \(z\) 的贡献。
  第一种情况直接一次dfs统计：

void dfs1(int x,int fa,int from){
    cnt_col[c[x]]++;
    if(cnt_col[c[x]]==1)ans[from]+=sz[x];
    for(auto to:G[x])if(to!=fa&&!vis[to])dfs1(to,x,from);
    cnt_col[c[x]]--;
}

对第二种情况，假设某种颜色 \(c\) 在 \(x\to y\) 的路径上已经出现过了，统计 \(y\) 的答案时，贡献直接是 \(sz[x]-sz[p]\)，否则，应该对 \(p\) 以外的 \(x\) 子树进行统计，如果某个颜色在一个结点里第一次出现，则会产生其子树大小的贡献，记一个 \(path[c]\) 表示颜色 \(c\) 的贡献

void dfs2(int x,int fa,int def_val,bool tag){
    cnt_col[c[x]]++;
    if(cnt_col[c[x]]==1){
        sum_col+=def_val;
        sum_path-=path[c[x]];
    }
    ans[x]+=sum_path;
    if(tag)ans[x]+=sum_col;
    for(auto to:G[x])if(to!=fa&&!vis[to])
        dfs2(to,x,def_val,tag);

    if(cnt_col[c[x]]==1){
        sum_col-=def_val;
        sum_path+=path[c[x]];
    }
    cnt_col[c[x]]--;
}

![[asset/Pasted image 20240429220518.png]]

因为需要扣除掉 \(p\) 子树内的 \(path\) 数组，直接给数组做差不方便，这部分答案可以考虑对 \(x\) 的孩子正反做两次dfs

void calc_path(int x,int fa){
    cnt_col[c[x]]++;
    Q[++tot]=c[x];
    if(cnt_col[c[x]]==1){
        path[c[x]]+=sz[x];
        sum_path+=sz[x];
    }
    for(auto to:G[x])if(to!=fa&&!vis[to])calc_path(to,x);
    cnt_col[c[x]]--;
}
//...
auto work=[&](bool tag){
	clear();
	sum_col=sum_path=0;
	assert(cnt_col[c[x]]==0);
	for(auto to:G[x])if(!vis[to]){
		cnt_col[c[x]]++;
		sum_col=sz[x]-sz[to];
		dfs2(to,x,sz[x]-sz[to],tag);
		calc_path(to,x);
		sum_col=0;
		cnt_col[c[x]]--;
	}
};
work(0);
reverse(G[x].begin(),G[x].end());
work(1);

代码：

#include<bits/stdc++.h>
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define endl '\n'
#define fastio ios_base::sync_with_stdio(false);cin.tie(0);cout.tie(0)
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
const int INF=0x3f3f3f3f;
const int N=1e5+5;
int n,c[N],mx_c;
bool vis[N];
vector<vector<int>> G;

int cnt_col[N],path[N];
int tot,Q[N];

ll sum_path,sum_col;
ll ans[N];
int rt,sz[N],mx_rt;
void get_rt(int x,int fa,int sum){
    sz[x]=1;
    int mx=0;
    for(auto to:G[x])if(to!=fa&&!vis[to]){
        get_rt(to,x,sum);
        mx=max(mx,sz[to]);
        sz[x]+=sz[to];
    }
    mx=max(mx,sum-sz[x]);
    if(mx<mx_rt){
        mx_rt=mx;
        rt=x;
    }
}

void dfs1(int x,int fa,int from){
    cnt_col[c[x]]++;
    if(cnt_col[c[x]]==1)ans[from]+=sz[x];
    for(auto to:G[x])if(to!=fa&&!vis[to])dfs1(to,x,from);
    cnt_col[c[x]]--;
}
void dfs2(int x,int fa,int def_val,bool tag){
    cnt_col[c[x]]++;
    if(cnt_col[c[x]]==1){
        sum_col+=def_val;
        sum_path-=path[c[x]];
    }
    ans[x]+=sum_path;
    if(tag)ans[x]+=sum_col;
    for(auto to:G[x])if(to!=fa&&!vis[to])
        dfs2(to,x,def_val,tag);

    if(cnt_col[c[x]]==1){
        sum_col-=def_val;
        sum_path+=path[c[x]];
    }
    cnt_col[c[x]]--;
}
void calc_path(int x,int fa){
    cnt_col[c[x]]++;
    Q[++tot]=c[x];
    if(cnt_col[c[x]]==1){
        path[c[x]]+=sz[x];
        sum_path+=sz[x];
    }
    for(auto to:G[x])if(to!=fa&&!vis[to])calc_path(to,x);
    cnt_col[c[x]]--;
}
void dfz(int x,int sum=n){
    mx_rt=INF;
    get_rt(x,-1,sum);
    x=rt;
    get_rt(x,-1,sum);
    
    auto clear=[&](){
        rep(i,1,tot)path[Q[i]]=0;
        tot=0;
    };
    dfs1(x,-1,x);
    auto work=[&](bool tag){
        clear();
        sum_col=sum_path=0;
        assert(cnt_col[c[x]]==0);
        for(auto to:G[x])if(!vis[to]){
            cnt_col[c[x]]++;
            sum_col=sz[x]-sz[to];
            dfs2(to,x,sz[x]-sz[to],tag);
            calc_path(to,x);
            sum_col=0;
            cnt_col[c[x]]--;
        }
    };
    work(0);
    reverse(G[x].begin(),G[x].end());
    work(1);

    vis[x]=true;
    for(auto to:G[x])if(!vis[to])dfz(to,sz[to]);
}
int main(){
    fastio;
    cin>>n;
    rep(i,1,n){
        cin>>c[i];
        mx_c=max(mx_c,c[i]);
    }
    G=vector<vector<int>>(n+1);
    rep(i,1,n-1){
        int u,v;
        cin>>u>>v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfz(1);
    rep(i,1,n)cout<<ans[i]<<endl;
    return 0;
}