[SP1487 PT07J][BZOJ 1803] - Query on a tree III

洛谷传送门

BZOJ传送门

题目描述

You are given a node-labeled rooted tree with n nodes.
Define the query $(x, k)$: Find the node whose label is k-th largest in the subtree of the node x. Assume no two nodes have the same labels.

输入输出格式

输入格式：

The first line contains one integer $n (1 ≤ n ≤ 10^5)$. The next line contains n integers $li (0 ≤ li ≤ 10^9)$ which denotes the label of the i-th node.
Each line of the following $n - 1$ lines contains two integers $u, v$. They denote there is an edge between node $u$ and node $v$. Node $1$ is the root of the tree.
The next line contains one integer m $(1 ≤ m ≤ 10^4)$ which denotes the number of the queries. Each line of the next m contains two integers $x, k$. ($k$ ≤ the total node number in the subtree of $x$)

输出格式：

For each query (x, k), output the index of the node whose label is the k-th largest in the subtree of the node x.

输入输出样例

输入样例#1：

5
1 3 5 2 7
1 2
2 3
1 4
3 5
4
2 3
4 1
3 2
3 2

输出样例#1：

5
4
5
5

解题分析

每个节点子树的DFS序是连续的一段区间， 直接利用DFS序建立主席树差分即可。
代码如下:

#include <cstdio>
#include <cctype>
#include <algorithm>
#include <cstdlib>
#include <cstring>
#include <cmath>
#define R register
#define IN inline
#define gc getchar()
#define W while
#define MX 200005
#define ls tree[now].son[0]
#define rs tree[now].son[1]
template <class T>
IN void in(T &x)
{
    x = 0; R char c = gc;
    W (!isdigit(c)) c = gc;
    W (isdigit(c))
    x = (x << 1) + (x << 3) + c - 48, c = gc;
}
int dot, q, cnt, eul, cot, dif;
int head[MX], root[MX], fir[MX], out[MX], dat[MX], buf[MX], ind[MX];
struct Edge
{
    int to, nex;
}edge[MX << 1];
IN void addedge(const int &from, const int &to)
{
    edge[++cnt] = {to, head[from]};
    head[from] = cnt;
}
struct Node
{
    int son[2], sum, id;
}tree[MX << 4];
namespace PT
{
    IN int getid(const int &now)
    {return std::lower_bound(buf + 1, buf + 1 + dif, now) - buf;}
    void DFS(const int &now, const int &fa)//建空树
    {
        fir[now] = cot++;//预处理开始位置， 方便查询
        ind[cot] = now;
        for (R int i = head[now]; i; i = edge[i].nex)
        {
            if(edge[i].to == fa) continue;
            DFS(edge[i].to, now);
        }
        out[now] = cot;//结束位置
    }
    void build(int &now, const int &lef, const int &rig)
    {//建立空树
        now = ++cot;
        if(lef == rig) return;
        int mid = (lef + rig) >> 1;
        build(ls, lef, mid);
        build(rs, mid + 1, rig);
    }
    void insert(const int &pre, int &now, const int &lef, const int &rig, const int &tar, const int &id)
    {//插入主席树节点
        tree[now = ++cot] = tree[pre];
        ++tree[now].sum;
        if(lef == rig) return tree[now].id = id, void();
        int mid = lef + rig >> 1;
        if(tar <= mid) insert(tree[pre].son[0], ls, lef, mid, tar, id);
        else insert(tree[pre].son[1], rs, mid + 1, rig, tar, id);
    }
    IN int query(R int x, R int y, const int &lef, const int &rig, const int &kth)
    {
        if(lef == rig) return tree[y].id;
        int rem = tree[tree[y].son[0]].sum - tree[tree[x].son[0]].sum;
        int mid = lef + rig >> 1;
        if(rem < kth) return query(tree[x].son[1], tree[y].son[1], mid + 1, rig, kth - rem);
        else return query(tree[x].son[0], tree[y].son[0], lef, mid, kth);
    }
}
int main(void)
{
    int a, b;
    in(dot);
    for (R int i = 1; i <= dot; ++i) in(dat[i]), buf[i] = dat[i];
    for (R int i = 1; i < dot; ++i)
    in(a), in(b), addedge(a, b), addedge(b, a);
    PT::DFS(1, 0);
    in(q);
    std::sort(buf + 1, buf + 1 + dot);
    dif = std::unique(buf + 1, buf + 1 + dot) - buf - 1;
    PT::build(root[0], 1, dif);
    for (R int i = 1; i <= dot; ++i)
    PT::insert(root[i - 1], root[i], 1, dif, PT::getid(dat[ind[i]]), ind[i]);
    W (q--)
    {
        in(a), in(b);
        printf("%d\n", PT::query(root[fir[a]], root[out[a]], 1, dif, b));
    }
}