洛谷2633 Count on a tree 题解（主席树）

题目大意

给定一个\(n\)个点的数，有\(m\)个询问，每次寻味两点之间最短距离上第\(k\)小的点权

思路

树上主席树。

学习了序列上的主席树之后，会发现主席树其实是一种前缀操作（在每个结点都建立一颗线段树）。因此这题我们也可以这样做，我们定义\(T[u]\)表示从根节点到到第\(u\)号节点的前缀主席树

所以包含\(u\)到\(v\)上的所有信息的就是：\(T[u]+T[v]-T[lca]-T[fa[lca]]\)（前缀和与差分是逆操作，所以可以用差分还原）

那么明白了这些后，代码就很显然了

#include <bits/stdc++.h>
using namespace std ;
const int MAXN = 1e5 + 5 ;
int n , m , tot , L ;
int a[ MAXN ] , b[ MAXN ] ;
int dep[ MAXN ] , fa[ MAXN ][ 20 ] ;
struct Node {
    int next , to ;
} edge[ MAXN << 1 ] ;
int head[ MAXN ] , cnt ;
struct Tree {
    int l , r , sum ;
} T[ MAXN << 5 ] ;
int rt[ MAXN ] ;
inline int read () {
    int tot = 0 , f = 1 ; char c = getchar () ;
    while ( c < '0' || c > '9' ) { if ( c == '-' ) f = -1 ; c = getchar () ; }
    while ( c >= '0' && c <= '9' ) { tot = tot * 10 + c - '0' ; c = getchar () ; }
    return tot * f ;
}
inline int id ( int x ) { return lower_bound ( b + 1 , b + 1 + L , a[ x ] ) - b ; } 
inline void add ( int x , int y ) {
    edge[ ++ cnt ].next = head[ x ] ;
    edge[ cnt ].to = y ;
    head[ x ] = cnt ;
}
inline void build ( Tree &u , int l , int r ) {
    u.sum = 0 ;
    if ( l == r ) return ;
    int mid = ( l + r ) >> 1 ;
    build ( T[ u.l = ++ tot ] , l , mid ) ;
    build ( T[ u.r = ++ tot ] , mid + 1 , r ) ;
}
inline void modify ( Tree &u , Tree ff , int l , int r , int p ) {
    u.sum = ff.sum + 1 ;
    if ( l == r ) return ;
    int mid = ( l + r ) >> 1 ;
    if ( p <= mid ) modify ( T[ u.l = ++ tot ] , T[ ff.l ] , l , mid , p ) , u.r = ff.r ;
    else modify ( T[ u.r = ++ tot ] , T[ ff.r ] , mid + 1 , r , p ) , u.l = ff.l ;
}
inline int query ( Tree x , Tree y , Tree z , Tree f , int l , int r , int k ) {
    if ( l == r ) return l ;
    int res = T[ x.l ].sum + T[ y.l ].sum - T[ z.l ].sum - T[ f.l ].sum ;
    int mid = ( l + r ) >> 1 ;
    if ( res >= k ) return query ( T[ x.l ] , T[ y.l ] , T[ z.l ] , T[ f.l ] , l , mid , k ) ;
    else return query ( T[ x.r ] , T[ y.r ] , T[ z.r ] , T[ f.r ] , mid + 1 , r , k - res ) ; // warning!!
}
// 以上为主席树基本操作
inline void dfs ( int u , int father ) {
    modify ( T[ rt[ u ] = ++ tot ] , T[ rt[ father ] ] , 1 , L , id ( u ) ) ; //建主席树
    dep[ u ] = dep[ father ] + 1 ;
    fa[ u ][ 0 ] = father ;
    for ( int i = 1 ; i <= 18 ; i ++ )
        fa[ u ][ i ] = fa[ fa[ u ][ i - 1 ] ][ i - 1 ] ;
    for ( int i = head[ u ] ; i ; i = edge[ i ].next ) {
        int v = edge[ i ].to ;
        if ( v == father ) continue ;
        dfs ( v , u ) ;
    }
}
inline int Lca ( int x , int y ) {
    if ( dep[ x ] < dep[ y ] ) swap ( x , y ) ;
    if ( x == y ) return x ;
    for ( int i = 18 ; i >= 0 ; i -- ) {
        if ( dep[ fa[ x ][ i ] ] >= dep[ y ] ) x = fa[ x ][ i ] ;
    }
    if ( x == y ) return x ;
    for ( int i = 18 ; i >= 0 ; i -- ) {
        if ( fa[ x ][ i ] != fa[ y ][ i ] ) x = fa[ x ][ i ] , y = fa[ y ][ i ] ;
    }
    return fa[ x ][ 0 ] ;
}
//倍增LCA的基本操作
signed main () {
    n = read () ; m = read () ;
    for ( int i = 1 ; i <= n ; i ++ ) b[ i ] = a[ i ] = read () ;
    for ( int i = 1 ; i < n ; i ++ ) {
        int x = read () , y = read () ;
        add ( x , y ) ; add ( y , x ) ;
    }
    sort ( b + 1 , b + 1 + n ) ;
    L = unique ( b + 1 , b + 1 + n ) - b - 1 ;//离散化
    build ( T[ rt[ 0 ] = ++ tot ] , 1 , L ) ;
    dfs ( 1 , 0 ) ;
    int ans = 0 ;
    while ( m -- ) {
        int u = read () , v = read () , w = read () ;
        u ^= ans ;
        int lca = Lca ( u , v ) ;
        // cout << u << " " << v << " " << lca << endl ;
        ans = query ( T[ rt[ u ] ] , T[ rt[ v ] ] , T[ rt[ lca ] ] , T[ rt[ fa[ lca ][ 0 ] ] ] , 1 , L , w ) ;
        ans = b[ ans ] ;
        printf ( "%d\n" , ans ) ;
    }
    return 0 ;
}