JZOJ 5678. 【GDOI2018Day2模拟4.21】果树

Description

NiroBC 姐姐是个活泼的少女，她十分喜欢爬树，而她家门口正好有一棵果树，正好满足了她爬树的需求。
这颗果树有N个节点，节点标号 1…N。每个节点长着一个果子，第i个节点上的果子颜色为 Ci 。
NiroBC姐姐每天都要爬树，每天都要选择一条有趣的路径 (u,v) 来爬。
一条路径被称作有趣的，当且仅当这条路径上的果子的颜色互不相同。
(u,v) 和 (v,u) 被视作同一条路径。特殊地，(i,i) 也被视作一条路径，这条路径只含 i 一个果子，显然是有趣的。
NiroBC姐姐想知道这颗树上有多少条有趣的路径。

Input

第一行，一个整数 N，表示果树的节点个数。
第二行，N 个整数 C1 ,C2 ,…,CN ，表示 N 个果子的颜色。
接下来 N−1 行，每行两个整数 ui ,vi ，表示 ui和vi 之间有一条边
数据保证这N−1条边构成一棵树。

Output

一个整数，表示有趣的路径的数量。

Sample Input

输入1：

3
1 2 3
1 2
1 3

输入2：

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

Sample Output

输出1：

6

样例解释：

有 (1,1),(1,2),(1,3),(2,2),(2,3),(3,3) 共 6 条有趣的路径。

输出2：

8

样例解释：

有 (1,1),(1,3),(2,2),(2,4),(2,5),(3,3),(4,4),(5,5) 共 8 条有趣的路径。

Data Constraint

Solution

• 先贴题解：

• 我们将每种颜色的点两两看做一组限制，在 $N*N$ 平面上覆盖矩形。

• 找“$a_i$ 下面的 $p$ 点”用倍增可以解决。

• 设有某种颜色的点有 $t\ (t\le20)$ 个，那么复杂度就是 $O(\frac{N}{t}*t^2\ logN)=O(N*t\ logN)$ 。

• 之后把每个矩形拆成两组（分别代表加入和删除），按x坐标排序，用扫描线扫过。

• 维护一颗标记不下传的线段树，记录区间内空点的个数即可。

• 时间复杂度 $O(N*t\ logN)$ 。

Code

#include<cstdio>
#include<algorithm>
#include<vector>
#include<cmath>
#include<cctype>
using namespace std;
const int N=1e5+5,M=21;
struct data
{
    int x,y1,y2,p;
}a[N*M<<2];
struct segment
{
    int sum,c;
}f[N<<2];
int n,tot,qx,qy,qz;
long long ans;
int first[N],nex[N<<1],en[N<<1];
int dfn[N],size[N],dep[N],fa[N][M];
vector<int>c[N];
inline int read()
{
    int X=0,w=0; char ch=0;
    while(!isdigit(ch)) w|=ch=='-',ch=getchar();
    while(isdigit(ch)) X=(X<<1)+(X<<3)+(ch^48),ch=getchar();
    return w?-X:X;
}
inline void insert(int x,int y)
{
    nex[++tot]=first[x];
    first[x]=tot;
    en[tot]=y;
}
void dfs(int x)
{
    dfn[x]=++tot;
    dep[x]=dep[fa[x][0]]+1;
    size[x]=1;
    for(int i=first[x];i;i=nex[i])
        if(en[i]^fa[x][0])
        {
            fa[en[i]][0]=x;
            dfs(en[i]);
            size[x]+=size[en[i]];
        }
}
inline void cover(int x1,int y1,int x2,int y2)
{
    a[++tot]=(data){x1,y1,y2,1};
    a[++tot]=(data){x2+1,y1,y2,-1};
}
inline bool cmp(data x,data y)
{
    return x.x<y.x;
}
inline int getlca(int x,int y)
{
    for(int i=log2(dep[y]);i>=0;i--)
        if(dep[fa[y][i]]>dep[x]) y=fa[y][i];
    return y;
}
void make(int v,int l,int r)
{
    f[v].sum=r-l+1;
    if(l==r) return;
    int mid=l+r>>1;
    make(v<<1,l,mid);
    make(v<<1|1,mid+1,r);
}
void change(int v,int l,int r)
{
    if(qx<=l && r<=qy)
    {
        f[v].c+=qz;
        if(f[v].c) f[v].sum=0; else
            if(l==r) f[v].sum=1; else
                f[v].sum=f[v<<1].sum+f[v<<1|1].sum;
        return;
    }
    int mid=l+r>>1;
    if(qx<=mid) change(v<<1,l,mid);
    if(qy>mid) change(v<<1|1,mid+1,r);
    if(!f[v].c) f[v].sum=f[v<<1].sum+f[v<<1|1].sum; else f[v].sum=0;
}
int main()
{
    freopen("tree.in","r",stdin);
    freopen("tree.out","w",stdout);
    n=read();
    for(int i=1;i<=n;i++)
    {
        int x=read();
        c[x].push_back(i);
    }
    for(int i=1;i<n;i++)
    {
        int x=read(),y=read();
        insert(x,y);
        insert(y,x);
    }
    tot=0;
    dfs(1);
    for(int j=1;j<17;j++)
        for(int i=1;i<=n;i++) fa[i][j]=fa[fa[i][j-1]][j-1];
    tot=0;
    for(int i=1;i<=n;i++)
    {
        int num=c[i].size();
        if(num>1)
            for(int j=0;j<num-1;j++)
                for(int k=j+1;k<num;k++)
                {
                    int x=c[i][j],y=c[i][k];
                    if(dfn[x]>dfn[y]) swap(x,y);
                    if(dfn[x]<dfn[y] && dfn[y]<=dfn[x]+size[x]-1)
                    {
                        if(fa[y][0]==x)
                        {
                            cover(1,dfn[y],dfn[y]-1,dfn[y]+size[y]-1);
                            cover(dfn[y],1,dfn[y]+size[y]-1,dfn[y]-1);
                            if(dfn[y]+size[y]<=n)
                            {
                                cover(dfn[y]+size[y],dfn[y],n,dfn[y]+size[y]-1);
                                cover(dfn[y],dfn[y]+size[y],dfn[y]+size[y]-1,n);
                            }
                        }else
                        {
                            int z=getlca(x,y);
                            cover(1,dfn[y],dfn[z]-1,dfn[y]+size[y]-1);
                            cover(dfn[y],1,dfn[y]+size[y]-1,dfn[z]-1);
                            if(dfn[z]+size[z]<=n)
                            {
                                cover(dfn[z]+size[z],dfn[y],n,dfn[y]+size[y]-1);
                                cover(dfn[y],dfn[z]+size[z],dfn[y]+size[y]-1,n);
                            }
                        }
                    }else
                    {
                        cover(dfn[x],dfn[y],dfn[x]+size[x]-1,dfn[y]+size[y]-1);
                        cover(dfn[y],dfn[x],dfn[y]+size[y]-1,dfn[x]+size[x]-1);
                    }
                }
    }
    sort(a+1,a+1+tot,cmp);
    make(1,1,n);
    for(int i=1,j=1;i<=n;i++)
    {
        while(j<=tot && a[j].x==i)
        {
            qx=a[j].y1,qy=a[j].y2,qz=a[j].p;
            change(1,1,n);
            j++;
        }
        ans+=f[1].sum;
    }
    printf("%lld",(ans+n)>>1);
    return 0;
}