牛客寒假算法集訓營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;
}