SPOJ - COT Count on a tree（主席树 LCA）

​​link​​

题意：

u v k : 询问从节点 u 到 节点 v 的路径上的第k小的权值

思路：

• 树上第$k$小问题，考虑怎么从序列转到树上。
• 对于一个节点，先将父节点的信息复制过来，再在$root[i]$里插入$a[i]$
• 这样每个节点维护的信息就是本节点到根节点的信息
• 求$u->v$信息的时候，用$root[u]+root[v]−root[lca]−root[fa[lca]]$，点差分的式子
• 因为$u->v$路径的点是包括$lca$的。

代码：

const int maxn=100000+7,mod=1e9+7,inf=0x3f3f3f3f;

int n,m,a[maxn];
vector<int>num;
struct node{
  int l,r,cnt;
}tr[maxn*40];
vector<int>g[maxn];
int dp[maxn][23],dep[maxn],root[maxn],idx,lim;
int get_id(int x){
  return lower_bound(num.begin(), num.end(),x)-num.begin();
}

void insert(int pre,int &now,int l,int r,int pos){
  now=++idx;
  tr[now]=tr[pre];//全部复制
  tr[now].cnt++;//更新
  if(l==r) return ;
  int mid=(l+r)/2;
  if(mid>=pos) insert(tr[pre].l,tr[now].l,l,mid,pos);
  else insert(tr[pre].r,tr[now].r,mid+1,r,pos);
}

int query(int u,int v,int lca,int lca_fa,int l,int r,int k){
  if(l==r) return l;
  int suml=tr[tr[u].l].cnt+tr[tr[v].l].cnt-tr[tr[lca].l].cnt-tr[tr[lca_fa].l].cnt;
  int mid=(l+r)/2;
  if(suml>=k) return query(tr[u].l,tr[v].l,tr[lca].l,tr[lca_fa].l,l,mid,k);
  else return query(tr[u].r,tr[v].r,tr[lca].r,tr[lca_fa].r,mid+1,r,k-suml);
}

void dfs(int u,int fa,int d){
  insert(root[fa],root[u],0,num.size(),get_id(a[u]));//主席树插入
  dep[u]=d;dp[u][0]=fa;
  for(int i=1;i<=lim;i++)
    dp[u][i]=dp[dp[u][i-1]][i-1];
  for(auto t:g[u]){
    if(t==fa) continue;
    dfs(t,u,d+1);
  }
}
int LCA(int x,int y)//求lca
{
  if(dep[x]<dep[y])
    swap(x,y);
  for(int i=lim;i>=0;i--)
    if(dep[x]-dep[y]>=(1<<i))
      x=dp[x][i];
  if(x==y)
    return x;
  for(int i=lim;i>=0;i--)
    if(dp[x][i]!=dp[y][i])
    {
      x=dp[x][i];
      y=dp[y][i];
    }
  return dp[x][0];
}


int main(){
  n=read,m=read;
  lim=log2(n);
  rep(i,1,n) a[i]=read,num.push_back(a[i]);
  sort(num.begin(),num.end());
  num.erase(unique(num.begin(),num.end()),num.end());
  rep(i,1,n-1){
    int u=read,v=read;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs(1,0,0);//
  rep(i,1,m){
    int u=read,v=read,k=read;
    int t=LCA(u,v);
    printf("%d\n",num[query(root[u],root[v],root[t],root[dp[t][0]],0,num.size(),k)]);
  }
  return 0;
}