spoj COT2 - Count on a tree II 树上莫队

题目链接

http://codeforces.com/blog/entry/43230树上莫队从这里学的，  受益匪浅..

#include <iostream>

#include <vector>

#include <cstdio>

#include <cstring>

#include <algorithm>

#include <cmath>

#include <map>

#include <set>

#include <string>

#include <queue>

#include <stack>

#include <bitset>

using namespace std;

#define pb(x) push_back(x)

#define ll long long

#define mk(x, y) make_pair(x, y)

#define lson l, m, rt<<1

#define mem(a) memset(a, 0, sizeof(a))

#define rson m+1, r, rt<<1|1

#define mem1(a) memset(a, -1, sizeof(a))

#define mem2(a) memset(a, 0x3f, sizeof(a))

#define rep(i, n, a) for(int i = a; i<n; i++)

#define fi first

#define se second

typedef pair<int, int> pll;

const double PI = acos(-1.0);

const double eps = 1e-;

const int mod = 1e9+;

const int inf = ;

const int dir[][] = { {-, }, {, }, {, -}, {, } };

const int maxn = 4e4+;

int in[maxn], head[maxn], num, out[maxn], cnt, val[maxn], ans[], a[maxn];

int f[maxn], d[maxn], p[maxn][], n, m, vis[maxn], times[maxn], res, id[maxn*];

struct node

{

    int to, nextt;

}e[maxn*];

struct query

{

    int l, r, lca, id, block;

    bool operator < (query a) const

    {

        if(block == a.block)

            return r<a.r;

        return block<a.block;

    }

}q[];

void add(int u, int v) {

    e[num].to = v, e[num].nextt = head[u], head[u] = num++;

}

void init() {

    mem1(head);

    mem1(p);

}

void dfs(int u, int fa) {

    in[u] = ++cnt;

    id[cnt] = u;

    f[u] = fa;

    for(int i = head[u]; ~i; i = e[i].nextt) {

        int v = e[i].to;

        if(v == fa)

            continue;

        d[v] = d[u]+;

        dfs(v, u);

    }

    out[u] = ++cnt;

    id[cnt] = u;

}

void initLca() {

    int i, j;

    for(i = ; i<=n; i++) {

        p[i][] = f[i];

    }

    for(j = ; (<<j)<=n; j++) {

        for(i = ; i<=n; i++) {

            if(~p[i][j-])

                p[i][j] = p[p[i][j-]][j-];

        }

    }

}

int lca(int u, int v) {

    if(d[u]<d[v])

        swap(u, v);

    int i, j;

    for(i = ; (<<i)<=d[u]; i++)

        ;

    i--;

    for(j = i; j>=; j--)

        if(d[u]-(<<j)>=d[v])

            u = p[u][j];

    if(u == v)

        return v;

    for(j = i; j>=; j--) {

        if(p[u][j] != - && p[u][j] != p[v][j]) {

            u = p[u][j];

            v = p[v][j];

        }

    }

    return f[u];

}

void check(int x) {

    if(vis[x] && --times[val[x]]==) {

        res--;

    } else if(vis[x]== && times[val[x]]++ == ) {

        res++;

    }

    vis[x]^=;

}

void cal() {

    int L = q[].l, R = q[].l-;

    for(int i = ; i<m; i++) {

        while(L<q[i].l) {

            check(id[L++]);

        }

        while(L>q[i].l) {

            check(id[--L]);

        }

        while(R<q[i].r) {

            check(id[++R]);

        }

        while(R>q[i].r) {

            check(id[R--]);

        }

        if(q[i].lca != id[q[i].l] && q[i].lca != id[q[i].r]) {

            check(q[i].lca);

        }

        ans[q[i].id] = res;

        if(q[i].lca != id[q[i].l] && q[i].lca != id[q[i].r]) {

            check(q[i].lca);

        }

    }

}

int main()

{

    int u, v;

    cin>>n>>m;

    int BLOCK = sqrt(n);

    init();

    for(int i = ; i<=n; i++) {

        scanf("%d", &val[i]);

        a[i-] = val[i];

    }

    sort(a, a+n);

    int N = unique(a, a+n)-a;

    for(int i = ; i<=n; i++) {

        val[i] = lower_bound(a, a+N, val[i])-a+;

    }

    for(int i = ; i<n-; i++) {

        scanf("%d%d", &u, &v);

        add(u, v);

        add(v, u);

    }

    d[] = ;

    dfs(, -);

    initLca();

    for(int i = ; i<m; i++) {

        scanf("%d%d", &u, &v);

        if(in[u]>in[v])

            swap(u, v);

        q[i].lca = lca(u, v);

        if(q[i].lca == u) {

            q[i].l = in[u];

            q[i].r = in[v];

        } else {

            q[i].l = out[u];

            q[i].r = in[v];

        }

        q[i].block = q[i].l/BLOCK;

        q[i].id = i;

    }

    sort(q, q+m);

    cal();

    for(int i = ; i<m; i++) {

        printf("%d\n", ans[i]);

    }

    return ;

}