牛客寒假算法集训营1 小a与军团模拟器（启发式合并）

这道题通过dfs暴力统计初始军团数是常规操作，此题难点在于之后的修改操作，我们可以用vector记录下每种颜色存在的位置，然后枚举要更改的颜色位置，每个位置原来对于答案的贡献是和他相邻不同色的节点个数，我们减去原来的答案加上更改后的答案就得到了当前的答案。然后再暴力合并vector就可以解决问题，然而这种解法会T掉，我们考虑优化合并操作，我们能发现把x变成y和把y变成x得到的答案是相同的，所以我们在合并的时候以及更改颜色的时候都是把颜色少的改成颜色多的，并且维护每种颜色实际的vector是哪个，这样一来可以把常数降低好几倍，也就是所谓的启发式合并。维护每种颜色实际的vector方法和并查集比较像，比较容易实现。

代码如下

#include<bits/stdc++.h>
using namespace std;
const int maxn=2e5+5;
vector<int> g[maxn];
vector<int> num[maxn];
int v[maxn];
int col[maxn];
int cnt;
int mp[maxn];
void dfs(int x)
{
    v[x]=true;
    num[col[x]].push_back(x);
    for(int i=0;i<g[x].size();i++)
    {
        int now=g[x][i];
        if(!v[now])
        {
            if(col[now]&&col[now]!=col[x])
                cnt++;
            dfs(now);
        }
    }
}
void work(int &x,int &y)
{
    if(num[x].size()>num[y].size())
        swap(x,y);
    for(int i=0;i<num[x].size();i++)
    {
        int now=num[x][i];
        for(int j=0;j<g[now].size();j++)
        {
            int t=g[now][j];
            if(col[t]&&col[t]!=x)
                cnt--;
        }
    }
    for(int i=0;i<num[x].size();i++)
    {
        int now=num[x][i];
        col[now]=y;
        num[y].push_back(now);
    }
    for(int i=0;i<num[x].size();i++)
    {
        int now=num[x][i];
        for(int j=0;j<g[now].size();j++)
        {
            int t=g[now][j];
            if(col[t]&&col[t]!=y)
                cnt++;
        }
    }
    num[x].clear();
    x=-1;
}
int main()
{
    int n,m;
    cin>>n>>m;
    for(int i=1;i<=n;i++)
        cin>>col[i];
    for(int i=1;i<=n;i++)
        mp[i]=i;
    int x,y;
    for(int i=1;i<n;i++)
    {
        cin>>x>>y;
        g[x].push_back(y);
        g[y].push_back(x);
    }
    cnt=(col[1]==0?0:1);
    dfs(1);
    for(int i=1;i<=m;i++)
    {
        cin>>x>>y;
        if(mp[y]==-1)
            mp[x]=-1,mp[y]=x;
        else if(mp[x]!=-1&&mp[x]!=mp[y])
            work(mp[x],mp[y]);
        cout<<cnt<<endl;
    }
    return 0;
}