线段树 BZOJ3252 攻略

  1 #include<iostream>
  2 #include<cstdio>
  3 #include<cstring>
  4 #include<algorithm>
  5 using namespace std;
  6 #define inf 1000000000000000ll 
  7 int n,k,cnt,tot;
  8 long long ans;
  9 long long val[200010],sum[200010];
 10 int head[200010],getin[200010],getout[200010],oula[200010],fa[200010];
 11 bool check[200010];
 12 struct data{
 13     int next,to;
 14 }edge[400010];
 15 struct dt{
 16     long long mx,loc,tag;
 17 }segtree[800010]; 
 18 void add(int start,int end){
 19     edge[++cnt].next=head[start];
 20     edge[cnt].to=end;
 21     head[start]=cnt;
 22     fa[end]=start;
 23 }
 24 void dfs(int u){
 25     getin[u]=++tot;
 26     oula[tot]=u;
 27     for(int i=head[u];i;i=edge[i].next){
 28         sum[edge[i].to]=sum[u]+val[edge[i].to];
 29         dfs(edge[i].to);
 30     }
 31     getout[u]=tot;
 32 }
 33 void up(int pos){
 34     segtree[pos].mx=0;
 35     if(segtree[pos<<1].mx>segtree[pos].mx){
 36         segtree[pos].mx=segtree[pos<<1].mx;
 37         segtree[pos].loc=segtree[pos<<1].loc;
 38     }
 39     if(segtree[pos<<1|1].mx>segtree[pos].mx){
 40         segtree[pos].mx=segtree[pos<<1|1].mx;
 41         segtree[pos].loc=segtree[pos<<1|1].loc;       
 42     }
 43 }
 44 void build(int pos,int ll,int rr){
 45     if(ll==rr){
 46         segtree[pos].mx=sum[oula[ll]];
 47         segtree[pos].loc=ll;
 48         return;
 49     }
 50     int mid=(ll+rr)>>1;
 51     build(pos<<1,ll,mid);
 52     build(pos<<1|1,mid+1,rr);
 53     up(pos);
 54 }
 55 void update(int pos,long long dd){//!!!!!!!!!!!!!!!!!!!!!!!! long long
 56     segtree[pos].mx-=dd;
 57     segtree[pos].tag+=dd;
 58 }
 59 void down(int pos){
 60     if(!segtree[pos].tag) return;
 61     update(pos<<1,segtree[pos].tag);
 62     update(pos<<1|1,segtree[pos].tag);
 63     segtree[pos].tag=0;
 64     return;
 65 }
 66 void change(int pl,int pr,long long dd,int pos,int ll,int rr){
 67     int mid=(ll+rr)>>1;
 68     if(pl>rr||pr<ll) return;
 69     if(pl<=ll&&pr>=rr){
 70         update(pos,dd);
 71         return;
 72     }
 73     down(pos);
 74     if(pr<=mid) change(pl,pr,dd,pos<<1,ll,mid);
 75     else if(pl>mid) change(pl,pr,dd,pos<<1|1,mid+1,rr);
 76     else{
 77         change(pl,mid,dd,pos<<1,ll,mid);
 78         change(mid+1,pr,dd,pos<<1|1,mid+1,rr);
 79     }
 80     up(pos);
 81 }
 82 int main(){
 83     scanf("%d%d",&n,&k);
 84     for(int i=1;i<=n;i++) scanf("%lld",&val[i]);
 85     int x,y;
 86     for(int i=1;i<n;i++){
 87         scanf("%d%d",&x,&y);
 88         add(x,y);
 89     }
 90     sum[1]=val[1];
 91     dfs(1);
 92     build(1,1,n);
 93     for(int i=1;i<=k;i++){
 94         ans+=segtree[1].mx;
 95         for(int j=oula[segtree[1].loc];!check[j]&&j;j=fa[j]){
 96             check[j]=1;
 97             change(getin[j],getin[j],inf,1,1,n);//ll的极值？？？ 
 98             if(getin[j]<getout[j]) change(getin[j]+1,getout[j],val[j],1,1,n);
 99         }
100     }
101     printf("%lld",ans);
102     return 0;
103 }