【SPOJ10707】Count on a tree II-树上莫队算法

测试地址：Count on a tree II
题目大意：有一棵树，每个点有一个点权，多个询问，每次询问一条路径上有多少个不同的点权。
做法：本题需要用到树上莫队算法。
我们知道这题如果出在序列上就是莫队算法的裸题，而这题就是树上莫队算法的裸题。
我们知道要使用莫队算法，必须要将一条路径化为一个区间。从这个意义上来说，树上莫队算法的思想和树链剖分类似，都是将一棵树化为一个序列，但是树链剖分是将一条路径和若干个序列上的区间对应，而树上莫队算法需要将一条路径和一个序列上的区间对应。树上莫队算法的转化方式是这样的：DFS整棵树，在每个点入栈和出栈时，都将这个点的编号记下，按这样的顺序组成一个长为$2N$的序列。
令点$i$入栈和出栈的时刻为$in_i$和$out_i$，对于一条路径$(a,b)$，有以下几种情况：
一、$a$是$b$的祖先，那么这条路径和序列上的区间$[in_a,in_b]$对应。
二、$b$是$a$的祖先，那么这条路径和序列上的区间$[in_b,in_a]$对应。
三、$a$和$b$无祖孙关系，那么这条路径和序列上的区间$[out_a,in_b]$（或$[out_b,in_a]$，因为我们不知道$a,b$的先后顺序）加上$lca(a,b)$这个点对应。
有时一个点在区间内出现两次，这种情况下我们不应该计算这个点，因为它表示这个点入栈又出栈了。上面第一、二种情况应该还好理解，第三种情况就有点让人头疼，一个区间加一个点是什么操作？实际上，由于$lca(a,b)$一定早于点$a,b$入栈时入栈，又一定晚于点$a,b$出栈时出栈，所以上面的区间并没有包含$lca(a,b)$，计算信息时要加上这个点。同时上述的区间定义可以被证明满足每个要算入答案的点都出现一次且仅一次，这里就不证了。
做完上述事情之后，就是一个裸的莫队算法了，时间复杂度$O(N\sqrt N)$。
注意点权是在$[0,2^{31}-1]$的范围，需要离散化。
以下是本人代码：

#include <bits/stdc++.h>
using namespace std;
int n,m,first[40010]={0},tot=0,in[40010],out[40010],tim=0,a[40010],p[80010];
int blocklen,block[80010],fa[40010][21]={0},dep[40010]={0},ans[100010];
int cnt[40010]={0},sum=0;
bool cntp[40010]={0};
struct edge {int v,next;} e[100010];
struct forsort {int id,val;} f[40010];
struct Query
{
    int id,l,r,extra;
}q[100010];

bool cmpf(forsort a,forsort b)
{
    return a.val<b.val;
}

bool cmpq(Query a,Query b)
{
    if (block[a.l]!=block[b.l]) return block[a.l]<block[b.l];
    else return a.r<b.r;
}

void insert(int a,int b)
{
    e[++tot].v=b,e[tot].next=first[a],first[a]=tot;
}

void dfs(int v)
{
    in[v]=++tim,p[tim]=v;
    for(int i=first[v];i;i=e[i].next)
        if (e[i].v!=fa[v][0])
        {
            fa[e[i].v][0]=v;
            dep[e[i].v]=dep[v]+1;
            dfs(e[i].v);
        }
    out[v]=++tim,p[tim]=v;
}

int lca(int x,int y)
{
    if (dep[x]<dep[y]) swap(x,y);
    for(int i=20;i>=0;i--)
        if (dep[fa[x][i]]>=dep[y]) x=fa[x][i];
    if (x==y) return x;
    for(int i=20;i>=0;i--)
        if (fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
    return fa[x][0];
}

void expand(int x)
{
    cntp[p[x]]=!cntp[p[x]];
    if (cntp[p[x]])
    {
        cnt[a[p[x]]]++;
        if (cnt[a[p[x]]]==1) sum++;
    }
    else
    {
        cnt[a[p[x]]]--;
        if (cnt[a[p[x]]]==0) sum--;
    }
}

void Mo()
{
    int l=1,r=0;
    for(int i=1;i<=m;i++)
    {
        while (q[i].l<l) expand(--l);
        while (q[i].r>r) expand(++r);
        while (q[i].l>l) expand(l++);
        while (q[i].r<r) expand(r--);
        if (q[i].extra&&!cnt[a[q[i].extra]]) ans[q[i].id]=sum+1;
        else ans[q[i].id]=sum;
    }
}

int main()
{
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)
    {
        f[i].id=i;
        scanf("%d",&f[i].val);
    }
    sort(f+1,f+n+1,cmpf);
    int totf=0;
    for(int i=1;i<=n;i++)
    {
        if (i==1||f[i].val!=f[i-1].val) totf++;
        a[f[i].id]=totf;
    }

    for(int i=1;i<n;i++)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        insert(a,b),insert(b,a);
    }
    dep[1]=1;
    dfs(1);
    for(int i=1;i<=20;i++)
        for(int j=1;j<=n;j++)
            fa[j][i]=fa[fa[j][i-1]][i-1];

    blocklen=(int)sqrt(n<<1);
    for(int i=1;i<=(n<<1);i++)
        block[i]=i/blocklen;

    for(int i=1;i<=m;i++)
    {
        int a,b,g;
        q[i].id=i;
        scanf("%d%d",&a,&b);
        g=lca(a,b);
        if (g==a||g==b) q[i].l=in[a],q[i].r=in[b],q[i].extra=0;
        else
        {
            q[i].extra=g;
            if (out[a]<=in[b]) q[i].l=out[a],q[i].r=in[b];
            else q[i].l=out[b],q[i].r=in[a];
        }
        if (q[i].l>q[i].r) swap(q[i].l,q[i].r);
    }

    sort(q+1,q+m+1,cmpq);
    Mo();
    for(int i=1;i<=m;i++)
        printf("%d\n",ans[i]);

    return 0;
}