SPOJCOT2 Count on a tree II

分析

树上莫队裸题。
好博客

树剖的时候不能再次dfs重儿子。（好像是废话，但我因为这个问题调了三小时）

代码

#include<cstdlib>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<ctime>
#include<iostream>
#include<string>
#include<vector>
#include<list>
#include<deque>
#include<stack>
#include<queue>
#include<map>
#include<set>
#include<bitset>
#include<algorithm>
#include<complex>
#pragma GCC optimize ("O0")
using namespace std;
template<class T> inline T read(T&x){
    T data=0;
    int w=1;
    char ch=getchar();
    while(!isdigit(ch))
    {
        if(ch=='-')
            w=-1;
        ch=getchar();
    }
    while(isdigit(ch))
        data=10*data+ch-'0',ch=getchar();
    return x=data*w;
}
typedef long long ll;
const int INF=0x7fffffff;

const int MAXN=1e5+7;
int n,q;
map<int,int>M;
int a[MAXN],id;
int belong[MAXN],block;

struct Edge
{
    int nx,to;
}E[MAXN<<2];
int head[MAXN],ecnt;

void addedge(int x,int y)
{
    E[++ecnt].to=y;
    E[ecnt].nx=head[x],head[x]=ecnt;
}

int fa[MAXN],dep[MAXN],sz[MAXN],son[MAXN],top[MAXN];
int st[MAXN],ed[MAXN],pot[MAXN],clk;

void dfs1(int x,int f)
{
    fa[x]=f,sz[x]=1;
    st[x]=++clk,pot[clk]=x;
    dep[x]=dep[f]+1;
    for(int i=head[x];i;i=E[i].nx)
    {
        int y=E[i].to;
        if(y==f)
            continue;
        dfs1(y,x);
        sz[x]+=sz[y];
        if(sz[y]>sz[son[x]])
            son[x]=y;
    }
    ed[x]=++clk,pot[clk]=x;
}

void dfs2(int x,int topf)
{
    top[x]=topf;
    if(!son[x])
        return;
    dfs2(son[x],topf);
    for(int i=head[x];i;i=E[i].nx)
    {
        int y=E[i].to;
        if(y==fa[x]||y==son[x]) // edit 1:cannot dfs son[x] twice and wrongly
            continue;
        dfs2(y,y);
    }
}

int lca(int x,int y)
{
    while(top[x]!=top[y])
    {
        if(dep[top[x]]<dep[top[y]])
            swap(x,y);
        x=fa[top[x]];
    }
    return dep[x]<dep[y]?x:y;
}

struct Quiz
{
    int l,r,f,id;
    bool operator<(const Quiz&x)
    {
        if(belong[l]^belong[x.l])
            return l<x.l;
        else
            return r<x.r;
    }
}Q[MAXN];

int ql,qr,res;
bool used[MAXN];
int cnt[MAXN],ans[MAXN];

void add(int x)
{
    if(++cnt[x]==1)
        ++res;
}

void del(int x)
{
    if(--cnt[x]==0)
        --res;
}

void change(int x)
{
    used[x]?del(a[x]):add(a[x]);
    used[x]^=1;
}

int main()
{
//  freopen(".in","r",stdin);
//  freopen(".out","w",stdout);
    read(n);read(q);
    block=sqrt(n);
    for(int i=1;i<=n;++i)
    {
        read(a[i]);
        if(!M[a[i]])
            M[a[i]]=++id;
        a[i]=M[a[i]];
    }
    for(int i=1;i<=2*n;++i) // edit 1:belong should be inited to 2n
        belong[i]=(i-1)/block+1;
    for(int i=1;i<n;++i)
    {
        int x,y;
        read(x);read(y);
        addedge(x,y);
        addedge(y,x);
    }
    dfs1(1,0);
    dfs2(1,1);
    for(int i=1;i<=q;++i)
    {
        int x,y;
        read(x);read(y);
        if(st[x]>st[y])
            swap(x,y);
        int f=lca(x,y);
        if(f==x)
            Q[i].l=st[x],Q[i].r=st[y];
        else
            Q[i].l=ed[x],Q[i].r=st[y],Q[i].f=f;
        Q[i].id=i;
    }
    sort(Q+1,Q+q+1);
    ql=1,qr=0,res=0;
    for(int i=1;i<=q;++i)
    {
        while(ql>Q[i].l)
            change(pot[--ql]);
        while(qr<Q[i].r)
            change(pot[++qr]);
        while(ql<Q[i].l)
            change(pot[ql++]);
        while(qr>Q[i].r)
            change(pot[qr--]);
        if(Q[i].f)
            change(Q[i].f);
        ans[Q[i].id]=res;
        if(Q[i].f)
            change(Q[i].f);
    }
    for(int i=1;i<=q;++i)
        printf("%d\n",ans[i]);
//  fclose(stdin);
//  fclose(stdout);
    return 0;
}