SPOJ 10628. SPOJ COT Count on a tree 可持久化线段树

#include <cstdio>
#include <cstdlib>
#include <algorithm>
#include <queue>
#include <map>

using namespace std;

const int maxn = 100005;
int n, m, w[maxn];
int head[maxn], label;
struct Edge
{
    int v, next;
    Edge (int v = 0, int next = 0):
        v(v), next(next) {}
}e[maxn*2];
int depth[maxn], f[maxn][31], root[maxn], temp[maxn], s_num;
queue <int> q;
map <int, int> id;
struct Tree
{
    int sum[maxn*20], ls[maxn*20], rs[maxn*20], cnt;
    Tree ()
    {
        sum[0] = 0, cnt = 0;
    }
    void PushUp(int rt)
    {
        sum[rt] = sum[ls[rt]]+sum[rs[rt]];
    }
    void insert(int las_rt, int rt, int l, int r, int p, int d)
    {
        if (l == r)
        {
            sum[rt] = sum[las_rt]+d;
            return ;
        }
        int mid = (l+r)>>1;
        if (p <= mid)
        {
            ls[rt] = ++cnt, rs[rt] = rs[las_rt];
            insert(ls[las_rt], ls[rt], l, mid, p, d);
        }
        else
        {
            ls[rt] = ls[las_rt], rs[rt] = ++cnt;
            insert(rs[las_rt], rs[rt], mid+1, r, p, d);
        }
        PushUp(rt);
    }
    int query(int u_rt, int v_rt, int lca_rt, int lca_pos, int l, int r, int k)
    {
        if (l == r)
            return l;
        int mid = (l+r)>>1;
        int s_l = sum[ls[u_rt]]+sum[ls[v_rt]]-2*sum[ls[lca_rt]]+(lca_pos >= l && lca_pos <= mid);
        if (k <= s_l)
            return query(ls[u_rt], ls[v_rt], ls[lca_rt], lca_pos, l, mid, k);
        else
            return query(rs[u_rt], rs[v_rt], rs[lca_rt], lca_pos, mid+1, r, k-s_l);
    }
}T;

void ins(int u, int v)
{
    e[++label] = Edge(v, head[u]);
    head[u] = label;
}

void in()
{
    scanf("%d %d", &n, &m);
    for (int i = 1; i <= n; ++i)
        scanf("%d", &w[i]);
    for (int i = 1; i <= n; ++i)
        head[i] = -1;
    label = -1;
    for (int i = 1; i < n; ++i)
    {
        int u, v;
        scanf("%d %d", &u, &v);
        ins(u, v), ins(v, u);
    }
}

void Build_lca()
{
    for (int i = 0; i <= 30; ++i)
        for (int j = 1; j <= n; ++j)
            f[j][i] = -1;
    q.push(1);
    depth[1] = 0;
    f[1][0] = -1;
    while (!q.empty())
    {
        int u = q.front();
        for (int i = head[u]; i != -1; i = e[i].next)
        {
            int v = e[i].v;
            if (v == f[u][0])
                continue ;
            f[v][0] = u;
            depth[v] = depth[u]+1;
            q.push(v);
        }
        q.pop();
    }
    for (int i = 1; i <= 30; ++i)
        for (int j = 1; j <= n; ++j)
        {
            if (f[j][i-1] == -1)
                continue ;
            f[j][i] = f[f[j][i-1]][i-1];
        }    
}

void Hash_a()
{
    for (int i = 1; i <= n; ++i)
        temp[i] = w[i];
    sort(temp+1, temp+n+1);
    s_num = 0;
    for (int i = 1; i <= n; ++i)
        if (temp[i] != temp[i-1] || i == 1)
        {
            temp[++s_num] = temp[i];
            id[temp[i]] = s_num;
        }
}

void dfs(int u)
{
    root[u] = ++T.cnt;
    T.insert(u == 1 ? 0 : root[f[u][0]], root[u], 1, s_num, id[w[u]], 1);
    for (int i = head[u]; i != -1; i = e[i].next)
    {
        int v = e[i].v;
        if (v == f[u][0])
            continue ;
        dfs(v);
    }
}

void Build_tree()
{
    Hash_a();
    dfs(1);/*
    for (int i = 1; i <= n; ++i)
        printf("%d %d\n", i, T.sum[root[i]]);*/
}

void prepare()
{
    Build_lca();
    Build_tree();
}

int lca(int u, int v)
{
    if (depth[u] < depth[v])
        swap(u, v);
    for (int i = 30; i >= 0; --i)
    {
        if (f[u][i] == -1)
            continue ;
        if (depth[f[u][i]] >= depth[v])
        {
            u = f[u][i];
            if (depth[u] == depth[v])
                break ;
        }
    }
    if (u == v)
        return u;
    for (int i = 30; i >= 0; --i)
    {
        if (f[u][i] == -1)
            continue ;
        if (f[u][i] != f[v][i])
        {
            u = f[u][i];
            v = f[v][i];
        }
    }
    return f[u][0];
}

void work()
{
    prepare();
    while (m --)
    {
        int u, v, k;
        scanf("%d %d %d", &u, &v, &k);
        int t = lca(u, v);
        int pos = T.query(root[u], root[v], root[t], id[w[t]], 1, s_num, k);
        printf("%d\n", temp[pos]);
    }
}

int main()
{
    in();
    work();
    return 0;
}