题解 P2633 Count on a tree

题目描述

Link

给定一棵 \(n\) 个节点的树，有 \(m\) 个询问，每次询问 \(x,y\) 路径上的 \(k\) 小权值。

强制在线。

\(1 \leq n,m \leq 10^5 ,1 \leq a_i \leq 2^{31}-1\) 。

Solution

首先，请先做完 【模板】可持久化线段树 2（主席树） 。

发现这题只是把静态区间 \(k\) 小搬到了树上。

我们回顾一下序列上 \(k\) 小的做法：在权值线段树上二分，利用主席树可以保存历史信息的特点，要查询 \([l,r]\) 在左儿子内有几个数，只需要把 \([l,r]\) 变成 \([1,r] - [1 ,l - 1]\) 即可。

那么对于这题，我们同样可以把一条路径 \((x,y)\) 拆成四个部分：

\[s_x+s_y-s_{lca}-s_{fa_{lca}} \]

其中 \(s_x\) 表示 \(x\) 到根， \(lca\) 表示 \(x,y\) 的最近公共祖先， \(fa_{lca}\) 表示最近公共祖先的父节点。

于是，我们可以先对树进行一次 \(dfs\) ，当 \(dfs\) 到 \(now\) 时，就在 \(rt_{fa_{now}}\) 的基础上（父亲这个版本的主席树）加上 \(a_{now}\) 这一个权值。

对于每个询问，我们可以直接在四个版本的主席树上面一起二分，和序列上的做法大同小异。

当然，对于 \(a\) 序列，我们要进行离散化。

至于求 \(lca\) ，这里用的是树链剖分（倍增也可以，但是更慢），同时在第一遍 \(dfs\) 的时候修改了主席树。

代码如下：

#include <cstdio>
#include <cstring>
#include <cctype>
#include <algorithm>
#include <iostream>
using namespace std;
inline int read() {
    int num = 0 ,f = 1; char c = getchar();
    while (!isdigit(c)) f = c == '-' ? -1 : f ,c = getchar();
    while (isdigit(c)) num = (num << 1) + (num << 3) + (c ^ 48) ,c = getchar();
    return num * f;
}
const int N = 1e5 + 5 ,M = N * 20; //空间开够
struct Segment {
    struct node {
        int l ,r ,sum;
        node (int l = 0 ,int r = 0 ,int sum = 0) : l(l) ,r(r) ,sum(sum) {}
    }t[M]; int tot;
    Segment() : tot(0) {}
    inline void update(int now) {
        t[now].sum = t[t[now].l].sum + t[t[now].r].sum;
    }
    inline void modify(int &now ,int l ,int r ,int q ,int k ,int p) {
        now = ++tot; t[now] = t[p];
        if (l == r) {t[now].sum += k; return ;}
        int mid = (l + r) >> 1;
        if (q <= mid) modify(t[now].l ,l ,mid ,q ,k ,t[p].l);
        else modify(t[now].r ,mid + 1 ,r ,q ,k ,t[p].r);
        update(now);
    }
    inline int ask(int a ,int b ,int c ,int d) {
        return t[a].sum + t[b].sum - t[c].sum - t[d].sum;
    }
    inline int query(int p ,int q ,int l ,int r ,int k ,int a ,int b) {
        if (l == r) return l;
        int mid = (l + r) >> 1 ,res = ask(t[p].l ,t[q].l ,t[a].l ,t[b].l);
        //写一个 ask 仅仅是为了没那么长
        if (res >= k) return query(t[p].l ,t[q].l ,l ,mid ,k ,t[a].l ,t[b].l);
        return query(t[p].r ,t[q].r ,mid + 1 ,r ,k - res ,t[a].r ,t[b].r);
    }
}t;
struct Edge {
    int to ,next;
    Edge (int to = 0 ,int next = 0) :
        to(to) ,next(next) {}
}G[N << 1]; int head[N] ,cnt;
inline void add(int u ,int v) {
    G[++cnt] = Edge(v ,head[u]); head[u] = cnt;
    G[++cnt] = Edge(u ,head[v]); head[v] = cnt;
}
int nums[N] ,tot;
int dep[N] ,fa[N] ,size[N] ,son[N] ,a[N] ,rt[N] ,top[N];
inline void dfs1(int now) {
    dep[now] = dep[fa[now]] + 1;
    size[now] = 1;
    t.modify(rt[now] ,1 ,tot ,a[now] ,1 ,rt[fa[now]]);
    for (int i = head[now]; i ; i = G[i].next) {
        int v = G[i].to;
        if (v == fa[now]) continue;
        fa[v] = now;
        dfs1(v);
        size[now] += size[v];
        if (size[v] > size[son[now]]) son[now] = v;
    }
}
inline void dfs2(int now ,int t) {
    top[now] = t;
    if (!son[now]) return ;
    dfs2(son[now] ,t);
    for (int i = head[now]; i ; i = G[i].next) {
        int v = G[i].to;
        if (v == fa[now] || v == son[now]) continue;
        dfs2(v ,v);
    }
}
inline 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;
}
int n ,m ,lastans;
signed main() {
    n = read(); m = read();
    for (int i = 1; i <= n; i++) nums[i] = a[i] = read();
    sort(nums + 1 ,nums + n + 1);
    tot = unique(nums + 1 ,nums + n + 1) - nums - 1;
    for (int i = 1; i <= n; i++)
        a[i] = lower_bound(nums + 1 ,nums + tot + 1 ,a[i]) - nums;
    //记得离散化。
    for (int i = 1; i <= n - 1; i++) {
        int u = read() ,v = read();
        add(u ,v);
    }
    dfs1(1); dfs2(1 ,1);
    while (m--) {
        int x = read() ^ lastans ,y = read() ,k = read();
        //记得解密
        int lca = LCA(x ,y);
        lastans = nums[t.query(rt[x] ,rt[y] ,1 ,tot ,k ,rt[lca] ,rt[fa[lca]])];
        printf("%d\n" ,lastans);
    }
    return 0;
}