https://loj.ac/problem/6276#submit_code

NiroBC 姐姐是個活潑的少女，她十分喜歡爬樹，而她家門口正好有一棵果樹，正好滿足了她爬樹的需求。
這顆果樹有N 個節點，節點標號1……N。每個節點長著一個果子，第i 個節點上的果子顏色為Ci。
NiroBC 姐姐每天都要爬樹，每天都要選擇一條有趣的路徑(u,v) 來爬。
一條路徑被稱作有趣的，當且僅當這條路徑上的果子的顏色互不相同。
(u,v) 和(v,u) 被視作同一條路徑。特殊地，(i,i) 也被視作一條路徑，這條路徑只含i 一個果子，顯然是有趣的。
NiroBC 姐姐想知道這顆樹上有多少條有趣的路徑。

線段樹好題（當然也很毒）。

考慮u和v同色時的不合法路徑，分兩種情況討論：

1.u和v有一個不同於二者的lca

顯然不合法路徑的兩端一個在u的子樹中，一個在v的子樹中。

2.v是u的祖先。

顯然不合法路徑的兩端一個在u的子樹中，一個在v的子樹的補集（包括v）中。

同時我們用dfs序定義（u，v）為起點為u終點為v的路徑，這樣我們可以發現不合法路徑的集合恰好圍成了多個矩陣。

那麽就是掃描線求矩形的並了（不會的話可以看POJ1389：Area of Simple Polygons），當然這是不合法路徑，你需要取反後把對角線加上/2才行。

（PPPS：對角線的點顯然不會屬於任何一個矩形，且其統計的方法不能/2，故需要加上。）

#include<cmath>
#include<cstdio>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
typedef long long ll;
typedef double dl;
const int N=1e5+5;
const int B=17;
inline int read(){
    int X=0,w=0;char ch=0;
    
 
while(!isdigit(ch)){w|=ch==‘-‘;ch=getchar();}
    while(isdigit(ch))X=(X<<3)+(X<<1)+(ch^48),ch=getchar();
    return w?-X:X;
}
struct path{
    int to,nxt;
}e[N*2];
struct node{
    int x1,x2,y,w;
}edge[N*160];
vector<int>c[N];
int n,cnt,head[N],pos[N],tot,num;
int anc[N][B+4],dep[N],size[N];
ll tr[N*4],lazy[N*4];
inline bool cmp(node a,node b){
    return a.y<b.y;
}
inline void add(int u,int v){
    e[++cnt].to=v;e[cnt].nxt=head[u];head[u]=cnt;
}
inline int LCA(int i,int j){
    if(dep[i]<dep[j])swap(i,j);
    for(int k=B;k>=0;k--)
        if(dep[anc[i][k]]>=dep[j])i=anc[i][k];
    if(i==j)return i;
    for(int k=B;k>=0;k--)
        if(anc[i][k]!=anc[j][k])
            i=anc[i][k],j=anc[j][k];
    return anc[i][0];
}
void dfs(int u){
    pos[u]=++tot;size[u]=1;
    dep[u]=dep[anc[u][0]]+1;
    for(int i=1;i<=B;i++)
        anc[u][i]=anc[anc[u][i-1]][i-1];
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v!=anc[u][0]){
            anc[v][0]=u;
            dfs(v);
            size[u]+=size[v];
        }
    }
}
void ins(int a,int l,int r,int l1,int r1,int w){
    if(r1<l||l1>r||l1>r1)return;
    if(l1<=l&&r<=r1){
        lazy[a]+=w;
        if(lazy[a]>0)tr[a]=r-l+1;
        else if(l==r)tr[a]=0;
        else tr[a]=tr[a*2]+tr[a*2+1];
        return;
    }
    int mid=(l+r)>>1;
    ins(a*2,l,mid,l1,r1,w);ins(a*2+1,mid+1,r,l1,r1,w);
    if(!lazy[a])tr[a]=tr[a*2]+tr[a*2+1];
}
int main(){
    ll n=read();
    for(int i=1;i<=n;i++)c[read()].push_back(i);
    for(int i=1;i<n;i++){
        int u=read(),v=read();
        add(u,v);add(v,u);
    }
    dfs(1);
    for(int i=1;i<=n;i++){
        for(int j=0;j<c[i].size();j++){
            for(int l=j+1;l<c[i].size();l++){
                int u=c[i][j],v=c[i][l];
                int lca=LCA(u,v);
                if(lca!=u&&lca!=v){
                    int x1=pos[u],y1=pos[v];
                    int x2=x1+size[u]-1,y2=y1+size[v]-1;
                    edge[++num]=(node){x1-1,x2,y1,1};
                    edge[++num]=(node){x1-1,x2,y2+1,-1};
                    edge[++num]=(node){y1-1,y2,x1,1};
                    edge[++num]=(node){y1-1,y2,x2+1,-1};
                }else{
                    if(dep[u]>dep[v])swap(u,v);
                    int p=v;
                    for(int k=B;k>=0;k--)
                        if(dep[anc[p][k]]>=dep[u]+1)p=anc[p][k];
                    int x1=pos[p],y1=pos[v];
                    int x2=x1+size[p]-1,y2=y1+size[v]-1;
                    edge[++num]=(node){0,x1-1,y1,1};
                    edge[++num]=(node){0,x1-1,y2+1,-1};
                    edge[++num]=(node){y1-1,y2,1,1};
                    edge[++num]=(node){y1-1,y2,x1,-1};
                    
                    edge[++num]=(node){x2,n,y1,1};
                    edge[++num]=(node){x2,n,y2+1,-1};
                    edge[++num]=(node){y1-1,y2,x2+1,1};
                    edge[++num]=(node){y1-1,y2,n+1,-1};
                }
            }
        }
    }
    ll ans=0;
    sort(edge+1,edge+num+1,cmp);
    ins(1,1,n,edge[1].x1+1,edge[1].x2,edge[1].w);
    for(int i=2;i<=num;i++){
        ll h=edge[i].y-edge[i-1].y;
        ans+=h*tr[1];
        ins(1,1,n,edge[i].x1+1,edge[i].x2,edge[i].w);
    }
    printf("%lld\n",(n*(n+1)-ans)>>1);
    return 0;
}

+++++++++++++++++++++++++++++++++++++++++++

+本文作者：luyouqi233。 　　　　　　　　　　　　　　+

+歡迎訪問我的博客：http://www.cnblogs.com/luyouqi233/+

+++++++++++++++++++++++++++++++++++++++++++

LOJ6276：果樹——題解