【BZOJ-3252】攻略 DFS序 + 线段树 + 贪心

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cmath>
#include<cstring>
using namespace std;
int read()
{
    int x=0,f=1; char ch=getchar();
    while (ch<‘0‘ || ch>‘9‘) {if (ch==‘-‘) f=-1; ch=getchar();}
    while (ch>=‘0‘ && ch<=‘9‘) {x=x*10+ch-‘0‘; ch=getchar();}
    return x*f;
}
#define MAXN 200010
#define LL long long
int N,K;
bool visit[MAXN];
LL sumV[MAXN],val[MAXN];
struct EdgeNode{int next,to;}edge[MAXN<<1];
int head[MAXN],cnt=1,fa[MAXN]; 
void AddEdge(int u,int v) {cnt++; edge[cnt].next=head[u]; head[u]=cnt; edge[cnt].to=v;}
void InsertEdge(int u,int v) {fa[v]=u; AddEdge(u,v); AddEdge(v,u);}
int pl[MAXN],pr[MAXN],dfn,pre[MAXN];
void DFS(int now,int last)
{
    pl[now]=++dfn; pre[dfn]=now;
    for (int i=head[now]; i; i=edge[i].next)
        if (edge[i].to!=last)
            sumV[edge[i].to]=sumV[now]+val[edge[i].to],DFS(edge[i].to,now);
    pr[now]=dfn;
}
struct SegmentTreeNode{int l,r,maxp; LL maxx,tag;}tree[MAXN<<2];
int Maxp(SegmentTreeNode ls,SegmentTreeNode rs) {return ls.maxx>rs.maxx? ls.maxp:rs.maxp;}
void Update(int now) 
{
    tree[now].maxx=max(tree[now<<1].maxx,tree[now<<1|1].maxx);
    tree[now].maxp=Maxp(tree[now<<1],tree[now<<1|1]);
}
void BuildTree(int now,int l,int r)
{
    tree[now].l=l; tree[now].r=r;
    if (l==r) {tree[now].maxx=sumV[pre[l]]; tree[now].maxp=l; return;}
    int mid=(l+r)>>1;
    BuildTree(now<<1,l,mid); BuildTree(now<<1|1,mid+1,r);
    Update(now);
}
void PushDown(int now)
{
    if (tree[now].l==tree[now].r || !tree[now].tag) return;
    LL tag=tree[now].tag; tree[now].tag=0;
    tree[now<<1].maxx+=tag; tree[now<<1|1].maxx+=tag;
    tree[now<<1].tag+=tag; tree[now<<1|1].tag+=tag;
}
void Change(int now,int L,int R,int val)
{
    PushDown(now);
    int l=tree[now].l,r=tree[now].r;
    if (L<=l && R>=r) {tree[now].maxx+=val; tree[now].tag+=val; return;}
    int mid=(l+r)>>1;
    if (L<=mid) Change(now<<1,L,R,val);
    if (R>mid) Change(now<<1|1,L,R,val);
    Update(now);
}
int main()
{
    N=read(); K=read();
    for (int i=1; i<=N; i++) val[i]=read();
    for (int x,y,i=1; i<=N-1; i++) x=read(),y=read(),InsertEdge(x,y);
    sumV[1]=val[1]; DFS(1,0);
//    for (int i=1; i<=N; i++)
//        printf("ID=%d   [%d , %d] %d %d\n",pre[i],pl[pre[i]],pr[pre[i]],val[pre[i]],sumV[pre[i]]);
    LL ans=0;
    BuildTree(1,1,N);
    for (int i=1; i<=K; i++)
        {
//            printf("NOW=%d   :  %d , %d\n",i,tree[1].maxx,tree[1].maxp);
            if (tree[1].maxx>0) ans+=tree[1].maxx; else break;
            for (int x=pre[tree[1].maxp]; !visit[x]&&x; x=fa[x])
                visit[x]=1,Change(1,pl[x],pr[x],-val[x]); 
        }
    printf("%lld\n",ans);
    return 0;
}