LOJ6276：果树——题解

https://loj.ac/problem/6276

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+;
const int B=;
inline int read(){
    int X=,w=;char ch=;
    while(!isdigit(ch)){w|=ch=='-';ch=getchar();}
    while(isdigit(ch))X=(X<<)+(X<<)+(ch^),ch=getchar();
    return w?-X:X;
}
struct path{
    int to,nxt;
}e[N*];
struct node{
    int x1,x2,y,w;
}edge[N*];
vector<int>c[N];
int n,cnt,head[N],pos[N],tot,num;
int anc[N][B+],dep[N],size[N];
ll tr[N*],lazy[N*];
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>=;k--)
        if(dep[anc[i][k]]>=dep[j])i=anc[i][k];
    if(i==j)return i;
    for(int k=B;k>=;k--)
        if(anc[i][k]!=anc[j][k])
            i=anc[i][k],j=anc[j][k];
    return anc[i][];
}
void dfs(int u){
    pos[u]=++tot;size[u]=;
    dep[u]=dep[anc[u][]]+;
    for(int i=;i<=B;i++)
        anc[u][i]=anc[anc[u][i-]][i-];
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v!=anc[u][]){
            anc[v][]=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]>)tr[a]=r-l+;
        else if(l==r)tr[a]=;
        else tr[a]=tr[a*]+tr[a*+];
        return;
    }
    int mid=(l+r)>>;
    ins(a*,l,mid,l1,r1,w);ins(a*+,mid+,r,l1,r1,w);
    if(!lazy[a])tr[a]=tr[a*]+tr[a*+];
}
int main(){
    ll n=read();
    for(int i=;i<=n;i++)c[read()].push_back(i);
    for(int i=;i<n;i++){
        int u=read(),v=read();
        add(u,v);add(v,u);
    }
    dfs();
    for(int i=;i<=n;i++){
        for(int j=;j<c[i].size();j++){
            for(int l=j+;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]-,y2=y1+size[v]-;
                    edge[++num]=(node){x1-,x2,y1,};
                    edge[++num]=(node){x1-,x2,y2+,-};
                    edge[++num]=(node){y1-,y2,x1,};
                    edge[++num]=(node){y1-,y2,x2+,-};
                }else{
                    if(dep[u]>dep[v])swap(u,v);
                    int p=v;
                    for(int k=B;k>=;k--)
                        if(dep[anc[p][k]]>=dep[u]+)p=anc[p][k];
                    int x1=pos[p],y1=pos[v];
                    int x2=x1+size[p]-,y2=y1+size[v]-;
                    edge[++num]=(node){,x1-,y1,};
                    edge[++num]=(node){,x1-,y2+,-};
                    edge[++num]=(node){y1-,y2,,};
                    edge[++num]=(node){y1-,y2,x1,-};                    edge[++num]=(node){x2,n,y1,};
                    edge[++num]=(node){x2,n,y2+,-};
                    edge[++num]=(node){y1-,y2,x2+,};
                    edge[++num]=(node){y1-,y2,n+,-};
                }
            }
        }
    }
    ll ans=;
    sort(edge+,edge+num+,cmp);
    ins(,,n,edge[].x1+,edge[].x2,edge[].w);
    for(int i=;i<=num;i++){
        ll h=edge[i].y-edge[i-].y;
        ans+=h*tr[];
        ins(,,n,edge[i].x1+,edge[i].x2,edge[i].w);
    }
    printf("%lld\n",(n*(n+)-ans)>>);
    return ;
}