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

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

分类：Mo's Algorithm

1. 题目链接

[spoj COT2- Count on a tree II]

2. 题意描述

给一棵 N 个节点的树，每个节点上有一个权值。然后 M 次询问，查询从 u→v 的路径上面的不同数的数目。
数据范围： N<=40000,M<=100000

3. 解题思路

其实是树上莫队的裸题，但是偶第一次写树上莫队，一堆bug。。。
首先，权值范围不确定，先对权值离散化一下。
然后直接跑树上莫队就好了。
树上路径上的莫队算法的步骤一般是：

1. 预处理出树的LCA，以及树上每个顶点进来的时间 st 和退出的时间 ed ，以及时间戳对应的顶点序号 seq ；
2. 询问 u→v(设st[u]≤st[v]) 的路径，无非是以下三种情况：
   
   • 当u==v时，路径上包含一个点；
   • 当u==lca(u,v)时，那么链上的点就是时间戳从 st[u] ~ st[v] 中的所有出现一次的点，即 seq 中下标从 st[u] ~ st[v] 且只出现一次的点。
   • 当u!=lca(u,v)时，那么链上的点就是时间戳从 ed[u] ~ st[v] 中所有出现一次的点+lca(u,v)

可以参考一下vfleaking大佬的讲解。

4. 实现代码

#include 
using namespace std;

#define FIN(x) freopen(#x".txt", "r", stdin)

typedef long long LL;
void debug() { cout << endl; }
template<typename T, typename ...R> void debug (T f, R ...r) { cout << "[" << f << "]"; debug (r...); }

const int MAXN = 40005;
const int MAXM = 100005;
const int DEEP = 30;

int n, m, block;
struct Edge {
    int v, next;
} edge[MAXN << 1];
int head[MAXN], tot, st[MAXN], ed[MAXN], tim, dep[MAXN], fa[MAXN][DEEP];
int seq[MAXN << 1];
int w[MAXN], f[MAXN], fsz;

void init_edge() {
    tim = tot = 0;
    memset(head, -1, sizeof(head));
}
inline void add_edge(int u, int v) {
    edge[tot] = Edge{v, head[u]};
    head[u] = tot ++;
}

void dfs(int u, int pre, int deep) {
    int v;
    dep[u] = deep; fa[u][0] = pre;
    st[u] = ++ tim; seq[tim] = u;
    for(int i = head[u]; ~i; i = edge[i].next) {
        v = edge[i].v;
        if(v == pre || v == u) continue;
        dfs(v, u, deep + 1);
    }
    ed[u] = ++ tim; seq[tim] = u;
}

int lca(int u, int v) {
    while(dep[u] != dep[v]) {
        if(dep[u] < dep[v]) swap(u, v);
        int d = dep[u] - dep[v];
        for(int i = 0; i < DEEP; ++i) {
            if(d >> i & 1) u = fa[u][i];
        }
    }
    if(u == v) return u;
    for(int i = DEEP - 1; i >= 0; --i) {
        if(fa[u][i] != fa[v][i]) {
            u = fa[u][i];
            v = fa[v][i];
        }
    }
    return fa[u][0];
}

struct Query {
    int L, R, bid, qid;
    int lca;
    Query() {}
    Query(int _L, int _R, int _lca, int _qid) : L(_L), R(_R), lca(_lca), qid(_qid) {
        bid = L / block;
    }
    bool operator < (const Query& e) {
        if(bid == e.bid) return R < e.R;
        return bid < e.bid;
    }
} qr[MAXM];
int Ans[MAXM];

struct Set {
    int vis[MAXN];
    int sz;
    void clear() { sz = 0; memset(vis, 0, sizeof(vis)); }
    void insert(int x) {
        if(!vis[x]) ++ sz;
        ++ vis[x];
    }
    void erase(int x) {
        -- vis[x];
        if(!vis[x]) -- sz;
    }
    int size() { return sz; }
} path_set;
bool in_path[MAXN];
inline void QModify(int x) {
    int u = seq[x];
    if(in_path[u]) path_set.erase(w[u]);
    else path_set.insert(w[u]);
    in_path[u] ^= 1;
}
int main() {
#ifdef ___LOCAL_WONZY___
    FIN(input);
#endif // ___LOCAL_WONZY___
    int u, v;
    scanf("%d %d", &n, &m);
    fsz = 0;
    for(int i = 1; i <= n; ++i) {
        scanf("%d", &w[i]);
        f[++ fsz] = w[i];
    }
    sort(f + 1, f + fsz + 1);
    fsz = unique(f + 1, f + fsz + 1) - f - 1;
    for(int i = 1; i <= n; ++i)  w[i] = lower_bound(f + 1, f + fsz + 1, w[i]) - f;
    init_edge();
    for(int i = 2; i <= n; ++i) {
        scanf("%d %d", &u, &v);
        add_edge(u, v);
        add_edge(v, u);
    }
    dfs(1, 1, 0);
    for(int i = 1; i < DEEP; i++) {
        for(int j = 1; j <= n; j++) {
            fa[j][i] = fa[fa[j][i - 1]][i - 1];
        }
    }
    block = sqrt(n * 2);
    int mm = 0;
    for(int i = 1; i <= m; ++i) {
        scanf("%d %d", &u, &v);
        if(u == v) {
            Ans[i] = 1;
            continue;
        }
        if(st[u] > st[v]) swap(u, v);
        int p = lca(u, v);
        if(p == u) {
            qr[++ mm] = Query(st[u], st[v], p, i);
        } else {
            qr[++ mm] = Query(ed[u], st[v], p, i);
        }
    }
    sort(qr + 1, qr + mm + 1);
    int L = 1, R = 0;
    memset(in_path, 0, sizeof(in_path));
    path_set.clear();

    for(int i = 1; i <= mm; ++i) {
        while(L > qr[i].L) QModify(--L);    //ins
        while(R < qr[i].R) QModify(++R);   //ins
        while(L < qr[i].L) QModify(L++);   //del
        while(R > qr[i].R) QModify(R--);   //del
        if(qr[i].lca != seq[qr[i].L] && qr[i].lca != seq[qr[i].R]) QModify(st[qr[i].lca]);  // ins
        Ans[qr[i].qid] = path_set.size();
        if(qr[i].lca != seq[qr[i].L] && qr[i].lca != seq[qr[i].R]) QModify(st[qr[i].lca]);  // del
    }
    for(int i = 1; i <= m; ++i) printf("%d\n", Ans[i]);
#ifdef ___LOCAL_WONZY___
    cout << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC * 1000 << "ms." << endl;
#endif // ___LOCAL_WONZY___
    return 0;
}