不同 ron clear 路徑 貢獻 utc ++i max long long

題目描述

lrb有一棵樹，樹的每個節點有個顏色。給一個長度為n的顏色序列，定義s(i,j) 為i 到j 的顏色數量。以及

現在他想讓你求出所有的sum[i]

輸入輸出格式

輸入格式：

第一行為一個整數n，表示樹節點的數量

第二行為n個整數，分別表示n個節點的顏色c[1],c[2]……c[n]

接下來n-1行，每行為兩個整數x,y，表示x和y之間有一條邊

輸出格式：

輸出n行，第i行為sum[i]

輸入輸出樣例

輸入樣例#1：

5
1 2 3 2 3
1 2
2 3
2 4
1 5

輸出樣例#1：

10
9
11
9
12

說明

sum[1]=s(1,1)+s(1,2)+s(1,3)+s(1,4)+s(1,5)=1+2+3+2+2=10
sum[2]=s(2,1)+s(2,2)+s(2,3)+s(2,4)+s(2,5)=2+1+2+1+3=9
sum[3]=s(3,1)+s(3,2)+s(3,3)+s(3,4)+s(3,5)=3+2+1+2+3=11
sum[4]=s(4,1)+s(4,2)+s(4,3)+s(4,4)+s(4,5)=2+1+2+1+3=9
sum[5]=s(5,1)+s(5,2)+s(5,3)+s(5,4)+s(5,5)=2+3+3+3+1=12
 

對於40%的數據，n<=2000

對於100%的數據，1<=n,c[i]<=10^5

題解

用的點分治做
這道題我看題解都看了很久啊
nyg和zlt用虛樹做的，於是我就只能默默地一個人看點分的做法
首先看找到分治中心後看怎麽算根的答案，因為點對中一個點一定是根，所以就是求每個點到根的路徑上有多少個不同的顏色；然後把計算答案的方法變一下，把統計點的貢獻變成統計顏色的貢獻
那麽如果一個點的顏色是在這個點到根的路徑上第一次出現，那麽這個顏色就可以對答案貢獻當前點的size大小貢獻（因為點對中另一個點只要是這個點的子樹中的點，那麽由於會經過當前點，而這個點的顏色又是第一次出現，那麽肯定每個點對的貢獻都會加1，那麽對於這個顏色來說，就會加size貢獻）

\(colvl[x]\) 代表顏色 \(x\) 的貢獻，\(allval\) 就是統計 \(colvl\) 的和
那麽根的答案就直接為 \(allval\)
由於分治算的答案都是必經過根的，所以我們接著會發現開始我們統計的 \(colvl[i]\) 同樣使用於除根外的其它點 \(x\) ，但要保證 \(x\) 到根的路徑上不能出現 \(i\) 的顏色，並且同一子樹中的點不能對其顏色有貢獻，這不就是點分不去重的搞法嗎
然後就好了
每次找完root後，dfs一遍算 \(colvl\) 和 \(allval\)，然後把根的貢獻搞出來
然後枚舉每一個子樹，dfs一遍把當前的子樹的貢獻去掉，再dfs一遍把去掉貢獻的子樹的答案算一下，最後dfs一遍把去掉的貢獻加回來
點分治就做完了

#include<bits/stdc++.h>
#define ui unsigned int
#define ll long long
#define db double
#define ld long double
#define ull unsigned long long
const int MAXN=100000+10,inf=0x3f3f3f3f;
int n,col[MAXN],e,to[MAXN<<1],nex[MAXN<<1],beg[MAXN],Msonsize[MAXN],size[MAXN],root,colnt[MAXN],finish[MAXN];
ll allval,colvl[MAXN],ans[MAXN];
template<typename T> inline void read(T &x)
{
    T data=0,w=1;
    char ch=0;
    while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();
    if(ch=='-')w=-1,ch=getchar();
    while(ch>='0'&&ch<='9')data=((T)data<<3)+((T)data<<1)+(ch^'0'),ch=getchar();
    x=data*w;
}
template<typename T> inline void write(T x,char ch='\0')
{
    if(x<0)putchar('-'),x=-x;
    if(x>9)write(x/10);
    putchar(x%10+'0');
    if(ch!='\0')putchar(ch);
}
template<typename T> inline void chkmin(T &x,T y){x=(y<x?y:x);}
template<typename T> inline void chkmax(T &x,T y){x=(y>x?y:x);}
template<typename T> inline T min(T x,T y){return x<y?x:y;}
template<typename T> inline T max(T x,T y){return x>y?x:y;}
inline void insert(int x,int y)
{
    to[++e]=y;
    nex[e]=beg[x];
    beg[x]=e;
}
inline void getroot(int x,int f,int total)
{
    Msonsize[x]=0;size[x]=1;
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==f||finish[to[i]])continue;
        else
        {
            getroot(to[i],x,total);
            size[x]+=size[to[i]];
            chkmax(Msonsize[x],size[to[i]]);
        }
    chkmax(Msonsize[x],total-size[x]);
    if(Msonsize[x]<Msonsize[root])root=x;
}
inline void dfs1(int x,int f)
{
    colnt[col[x]]++;
    size[x]=1;
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==f||finish[to[i]])continue;
        else dfs1(to[i],x),size[x]+=size[to[i]];
    if(colnt[col[x]]==1)
    {
        allval+=size[x];
        colvl[col[x]]+=size[x];
    }
    colnt[col[x]]--;
}
inline void dfs2(int x,int f,int k)
{
    colnt[col[x]]++;
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==f||finish[to[i]])continue;
        else dfs2(to[i],x,k);
    if(colnt[col[x]]==1)
    {
        allval+=k*size[x];
        colvl[col[x]]+=k*size[x];
    }
    colnt[col[x]]--;
}
inline void dfs3(int x,int f,int other,int colnm)
{
    colnt[col[x]]++;
    if(colnt[col[x]]==1)allval-=colvl[col[x]],colnm++;
    ans[x]+=(ll)allval+(ll)colnm*other;
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==f||finish[to[i]])continue;
        else dfs3(to[i],x,other,colnm);
    if(colnt[col[x]]==1)allval+=colvl[col[x]],colnm--;
    colnt[col[x]]--;
}
inline void clear(int x,int f)
{
    colnt[col[x]]=colvl[col[x]]=0;
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==f||finish[to[i]])continue;
        else clear(to[i],x);
}
inline void calc(int x)
{
    allval=0;
    dfs1(x,0);
    ans[x]+=allval;
    for(register int i=beg[x];i;i=nex[i])
        if(!finish[to[i]])
        {
            colnt[col[x]]++;
            allval-=size[to[i]];
            colvl[col[x]]-=size[to[i]];
            dfs2(to[i],x,-1);
            colnt[col[x]]--;
            dfs3(to[i],x,size[x]-size[to[i]],0);    
            colnt[col[x]]++;
            allval+=size[to[i]];
            colvl[col[x]]+=size[to[i]];
            dfs2(to[i],x,1);
            colnt[col[x]]--;
        }
    clear(x,0);
}
inline void solve(int x)
{
    calc(x);
    finish[x]=1;
    for(register int i=beg[x];i;i=nex[i])
        if(!finish[to[i]])
        {
            root=0;
            getroot(to[i],x,size[to[i]]);
            solve(root);
        }
}
int main()
{
    read(n);
    for(register int i=1;i<=n;++i)read(col[i]);
    for(register int i=1;i<n;++i)
    {
        int u,v;
        read(u);read(v);
        insert(u,v);insert(v,u);
    }
    Msonsize[0]=inf;
    getroot(1,0,n);
    solve(root);
    for(register int i=1;i<=n;++i)write(ans[i],'\n');
    return 0;
}

【刷題】洛谷 P2664 樹上遊戲