SP10628 COT - Count on a tree 树链第k小, dfs顺序上主席树

题目

洛谷 P2633 (https://www.luogu.com.cn/problem/P2633)

题意

给你一棵有n个结点的树，节点编号为1~n。

每个节点都有一个权值。

要求执行以下操作：

U V K：求从节点u到节点v的第k小权值。

输入格式

第一行有两个整数n和m（n，m≤100000） 第二行有n个整数。 第i个整数表示第i个节点的权值。

接下来的n-1行中，每行包含两个整数u v，表示u和v之间有一条边。

接下来的m行，每行包含三个整数U V K，进行一次操作。

输出格式

对于每个操作，输出结果。

输入

8 5
105 2 9 3 8 5 7 7
1 2
1 3
1 4
3 5
3 6
3 7
4 8
2 5 1
2 5 2
2 5 3
2 5 4
7 8 2

输出

2
8
9
105
7

思路

我们考虑在序列这个题应该怎么搞，就是主席树，主席就是一个前缀和，那么查询[L, R]区间
就是root[R]-root[L-1]。
在树上我们在每个节点建立线段树，继承父节点，那么查询[L, R]就是\(root[L]+root[R]-root[LCA]-root[fa[LCA]]\)
然后和在序列上是一样的。

#include<bits/stdc++.h>
#define LL long long
#define mid (l+r>>1)
using namespace std;

const int N=2e5+7;
const LL INF=1e18;

struct TreeLCA {
    struct zzz {
        LL t, w, nex;
    } e[500010 << 1];
    int head[500010], tot;
    void init(int n){
        tot=0;
        memset(head, 0, sizeof(head));
        for(int i = 1; i <= n; ++i){
            lg[i] = lg[i-1] + (1 << lg[i-1] == i);
        }
    }
    void add(LL x, LL y, LL w=1) {
        e[++tot].t = y;
        e[tot].w = w;
        e[tot].nex = head[x];
        head[x] = tot;
    }
    LL depth[300005], g[300005], fa[300005][22], lg[300005];
    void dfs(LL now, LL fath) {
        fa[now][0] = fath;
        depth[now] = depth[fath] + 1;
        for(LL i = 1; i <= lg[depth[now]]; ++i)
            fa[now][i] = fa[fa[now][i-1]][i-1];
        for(LL i = head[now]; i; i = e[i].nex)
            if(e[i].t != fath)
                g[e[i].t]=g[now]+e[i].w, dfs(e[i].t, now);
    }
    LL LCA(LL x, LL y) {
        if(depth[x] < depth[y])
            swap(x, y);
        while(depth[x] > depth[y])
            x = fa[x][lg[depth[x]-depth[y]] - 1];
        if(x == y)
            return x;
        for(LL k = lg[depth[x]] - 1; k >= 0; --k)
            if(fa[x][k] != fa[y][k])
                x = fa[x][k], y = fa[y][k];
        return fa[x][0];
    }
    LL dis(LL x, LL y) {
        return g[x]+g[y]-2*g[LCA(x, y)];
    }
}lca;

int root[N];
struct SegTree {
    LL sum[N*40], tot=0;
    int L[N*40], R[N*40];
    void init () {
        for(int i=0; i<=tot; i++){
            L[i]=R[i]=sum[i]=0;
        }
        tot=0;
    }
    int BT(int l, int r){
        int rt=++tot;
        sum[rt]=0;
        if(l<r){
            L[rt]=BT(l, mid);
            R[rt]=BT(mid+1, r);
        }
        return rt;
    }
    int add(int root, int l ,int r, int x, int val){//a[x]+=val
        int rt=++tot;
        L[rt]=L[root], R[rt]=R[root], sum[rt]=sum[root]+val;
        if(l<r){
            if(x<=mid) L[rt]=add(L[root], l, mid, x, val);
            else R[rt]=add(R[root], mid+1, r, x, val);
        }
        return rt;
    }
    int query(int A, int B, int LCA, int LCAfa, int l, int r, int k){//区间[x, y]的第k小
        if(l>=r) return l;//得到答案
        //cout<<lca.LCA(x, y)<<" "<<lca.fa[lca.LCA(x, y)][0]<<endl;
        int s=sum[L[A]]+sum[L[B]]-sum[L[LCA]]-sum[L[LCAfa]];

        if(s>=k) return query(L[A], L[B], L[LCA], L[LCAfa], l, mid, k);
        else return query(R[A], R[B], R[LCA], R[LCAfa], mid+1, r, k-s);
    }

}T;

struct LSH{//离散化
    int b[N];
    unordered_map<int, int> mp;//范围
    int lsh(int a[], int n){//得到离散化后不同元素的个数
        mp.clear();
        for(int i=1; i<=n; i++) b[i]=a[i];
        sort(b+1, b+n+1);
        int cnt=unique(b+1, b+n+1)-b-1;
        for(int i=1; i<=n; i++){
            int x=a[i];
            a[i]=lower_bound(b+1, b+cnt+1, a[i])-b;
            mp[a[i]]=x;
        }
        return cnt;
    }
    int id(int x){//得到原数
        return mp[x];
    }
}Lsh;

int w[N], cnt=0;

void DFS(int u, int fa){

    root[u]=T.add(root[fa], 1, cnt, w[u], 1);
    for(int i = lca.head[u]; i; i = lca.e[i].nex){
        int to=lca.e[i].t;
        if(to!=fa){
            DFS(to, u);
        }
    }
}

int main(){

    int n, m; scanf("%d%d", &n, &m);
    for(int i=1; i<=n; i++){
        scanf("%d", &w[i]);
    }
    cnt=Lsh.lsh(w, n);//离散化

    lca.init(n);
    T.init();

    for(int i=1; i<n; i++){
        int x, y; scanf("%d%d", &x, &y);
        lca.add(x, y);
        lca.add(y, x);
    }
    lca.dfs(1, 0);

    root[0]=T.BT(1, cnt);
    DFS(1, 0);

    while(m--){
        int u, v, k; scanf("%d%d%d", &u, &v, &k);
        int LCA=lca.LCA(u, v), LCAfa=lca.fa[LCA][0];
        
        int ans=T.query(root[u], root[v], root[LCA], root[LCAfa], 1, cnt, k);
        printf("%d\n", Lsh.id(ans));
    }

    return 0;
}

标签：Count,return,COT,int,LL,tree,fa,root,节点
来源： https://www.cnblogs.com/liweihang/p/13541074.html