C - Count on a tree（主席树+lca）

冷不防 2023-06-17 05:54 33阅读 0赞

题目链接:https://www.spoj.com/problems/COT/en/  
题意： 给你一颗树，问你u，v结点这条路径上第k大是多少。

分析：建主席树，儿子结点x的主席树由fa\[x\]得来，并更新。  
dist\[a\]+dist\[b\]-2\*dist\[lca(a,b)\]。  
这个题查询的地方写成了sum\[a\] 而不是sum\[lson\[a\]\] 导致一直出错，下次注意。

#include"string.h"
    #include"vector"
    #include"math.h"
    #include"queue"
    #include"algorithm"
    using namespace std;
    
    const int N = 100100;
    
    vector<int> Graph[N];
    int root[N],top,sum[N * 50];
    int lson[N * 50],rson[N * 50];
    int a[N],b[N],num[N],n,m;
    int len,d[N];
    int f[N][20],t;
    
    
    void bfs()
    {
        queue<int> q;
        q.push(1); d[1] = 1;
        while(q.size())
        {
            int x = q.front(); q.pop();
            for(int i = 0; i < Graph[x].size(); i ++)
            {
                int y = Graph[x][i];
                if(x == y || d[y]) continue;
                d[y] = d[x] + 1;
                f[y][0] = x;
                for(int j = 1; j <= t; j ++)
                    f[y][j] = f[f[y][j - 1]][j - 1];
                q.push(y);
            }
        }
    }
    int lca(int x,int y)
    {
        if(d[x] > d[y]) swap(x,y);
        for(int i = t; i >= 0; i --)
           if(d[f[y][i]] >= d[x]) y = f[y][i];
        if(x == y) return x;
        for(int i = t; i >= 0; i --)
            if(f[x][i] != f[y][i]) x = f[x][i],y = f[y][i];
        return f[x][0];
    }
    void Build(int &rt,int l,int r)
    {
        rt = ++ top;
        lson[rt] = 0; rson[rt] = 0;
        sum[rt] = 0;
        if(l == r)
            return ;
        int mid = (l + r) >> 1;
        Build(lson[rt],l,mid);
        Build(rson[rt],mid + 1,r);
    }
    void Update(int &ne,int last,int l,int r,int pos)
    {
        ne = ++ top;
        sum[ne] = sum[last] + 1;
        lson[ne] = lson[last]; rson[ne] = rson[last];
        if(l == r) return ;
        int mid = (l + r) >> 1;
        if(pos <= mid) Update(lson[ne],lson[last],l,mid,pos);
        else Update(rson[ne],rson[last],mid + 1,r,pos);
    }
    void dfs(int x,int fa)
    {
        Update(root[x],root[fa],1,len,num[x]);
        for(int i = 0; i < Graph[x].size(); i ++)
        {
            int y = Graph[x][i];
            if(y == fa) continue;
            dfs(y,x);
        }
    }
    int Query(int x,int y,int f,int fa,int k,int l,int r)
    {
        if(l == r) return l;
        int ans = sum[lson[x]] + sum[lson[y]] - sum[lson[f]] - sum[lson[fa]];
        int mid = (l + r) >> 1;
        if(k <= ans) return Query(lson[x],lson[y],lson[f],lson[fa],k,l,mid);
        else return Query(rson[x],rson[y],rson[f],rson[fa],k - ans,mid + 1,r);
    }
    void init()
    {
        memset(f,0,sizeof(f));
        top = 0;
        for(int i = 0; i <= n;i ++)
            Graph[i].clear();
    }
    int main()
    {
        while(~scanf("%d%d",&n,&m))
        {
            init();
            t = (int)(log(n)/log(2)) + 1;
            for(int i = 1; i <= n;i ++)
            {
                scanf("%d",&a[i]); b[i] = a[i];
            }
            sort(b + 1,b + 1 + n);
            len = unique(b + 1,b + 1 + n) - (b + 1);
            for(int i = 1; i <= n; i ++)
                num[i] = lower_bound(b + 1,b + len + 1,a[i]) - b;
            for(int i = 1; i < n; i ++)
            {
                int u,v; scanf("%d%d",&u,&v);
                Graph[u].push_back(v); Graph[v].push_back(u);
            }
            bfs();
            Build(root[0],1,len);
            dfs(1,0);
            /*for(int i = 1; i <= n; i ++)
                printf("rt = %d\n",root[i]);*/
            while(m --)
            {
                int x,y,k; scanf("%d%d%d",&x,&y,&k);
                int fa = lca(x,y);
                int ans = Query(root[x],root[y],root[fa],root[f[fa][0]],k,1,len);
              //  printf("ans = %d\n",ans);
                printf("%d\n",b[ans]);
            }
        }
    }