HDU 6162 Ch’s gift

做法同树上第K大

#include<stdio.h>
#include<string.h>
#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
const int MAXN =(int)1e5+10;
typedef long long ll;
int vl[MAXN],n,m,vsz;
vector<int>ve;
int root[MAXN],ls[MAXN*40],rs[MAXN*40],tot,dep[MAXN],f[MAXN][20];
ll sum[MAXN*40];
void insert(int l,int r,int &x,int y,int pos){
    x=++tot;ls[x]=ls[y],rs[x]=rs[y],sum[x]=sum[y]+ve[pos-1];
    if(l==r)return;
    int mid=l+r>>1;
    if(mid>=pos)insert(l,mid,ls[x],ls[y],pos);
    else insert(mid+1,r,rs[x],rs[y],pos);
}
ll query(int l,int r,int u,int v,int a,int b,int k){
    
    if(l==r)return sum[u]+sum[v]-sum[a]-sum[b];
    int mid=l+r>>1;
    if(k<=mid)return query(l,mid,ls[u],ls[v],ls[a],ls[b],k);
    else{
         return sum[ls[u]]+sum[ls[v]]-sum[ls[a]]-sum[ls[b]]+query(mid+1,r,rs[u],rs[v],rs[a],rs[b],k);
    }
}
int lca(int a,int b){
    int t,i;
    if(dep[a]<dep[b])t=a,a=b,b=t;
    for(i=16;i>=0;i--){
        if(dep[f[a][i]]>=dep[b])a=f[a][i];
    }
    if(a==b)return a;
    for(i=16;i>=0;i--){
        if(f[a][i]!=f[b][i]&&f[a][i]!=0&&f[b][i]!=0)a=f[a][i],b=f[b][i];
    }
    return f[a][0];
}
struct edge{
    int nxt,to;
}ed[MAXN<<1];
int head[MAXN],cnt;
void addedge(int u,int v){
    ed[cnt].to=v;
    ed[cnt].nxt=head[u];
    head[u]=cnt++;
}
void dfs(int u,int pre){
    insert(1,vsz,root[u],root[pre],vl[u]);
    for(int i=head[u];i!=-1;i=ed[i].nxt){ 
        int v=ed[i].to; 
        if(!dep[v]){            
          dep[v]=dep[u]+1;
          f[v][0]=u; 
          dfs(v,u);
        }
    }    
}
ll query(int u,int v,int k){
    if(k==0)return 0;
    int LCA=lca(u,v);
    return query(1,vsz,root[u],root[v],root[LCA],root[f[LCA][0]],k);
}
ll ans[MAXN];
int main()
{
    while(~scanf("%d%d",&n,&m)){
        ve.clear();
        for(int i=1;i<=n;i++)head[i]=-1,dep[i]=0,root[i]=0;
        for(int i=1;i<=n;i++)for(int j=1;j<=20;j++)f[i][j]=0;
        cnt=0;
        for(int i=1;i<=n;i++)scanf("%d",&vl[i]),ve.push_back(vl[i]);
        sort(ve.begin(),ve.end());
        ve.erase(unique(ve.begin(),ve.end()),ve.end());
        vsz=ve.size();
        for(int i=1;i<=n;i++)vl[i]=lower_bound(ve.begin(),ve.end(),vl[i])-ve.begin()+1;
        for(int i=1;i<n;i++){
            int u,v;
            scanf("%d%d",&u,&v);
            addedge(u,v);
            addedge(v,u);
        }
        dep[1]=1;
        dfs(1,0);
        for(int j=1;(1<<j)<=n;j++)for(int i=1;i<=n;i++)if(f[i][j-1]!=0)f[i][j]=f[f[i][j-1]][j-1];
        for(int i=1;i<=m;i++){
            int u,v,x,y;
            scanf("%d%d%d%d",&u,&v,&x,&y);
            y=upper_bound(ve.begin(),ve.end(),y)-ve.begin();
            x=lower_bound(ve.begin(),ve.end(),x)-ve.begin();
            ans[i]=query(u,v,y)-query(u,v,x);
        }
        for(int i=1;i<=m;i++){
            printf("%lld",ans[i]);
            if(i!=m)printf(" ");
            else printf("\n");
        }
        for(int i=0;i<=tot;i++)ls[i]=rs[i]=sum[tot]=0;
        tot=0;
    }
    return 0;
}