E71 树形DP+二分 P3523 [POI2011] DYN-Dynamite

视频链接：

P3523 [POI2011] DYN-Dynamite - 洛谷 | 计算机科学教育新生态

// 树形DP+二分 O(nlogn)
#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;
int read(){
  int x=0,f=1;char c=getchar();
  while(c>'9'||c<'0'){if(c=='-') f=-1;c=getchar();}
  while(c>='0'&&c<='9'){x=x*10+c-'0';c=getchar();}
  return x*f;
}

const int N=300005;
int idx,head[N],to[N<<1],ne[N<<1];
void add(int x,int y){
  to[++idx]=y;ne[idx]=head[x];head[x]=idx;
}
int n,m,mid,tot,b[N];
int f[N],g[N];

void dfs(int u,int fa){
  f[u]=-1e9;g[u]=1e9;
  for(int i=head[u];i;i=ne[i]){
    int v=to[i];
    if(v==fa) continue;
    dfs(v,u);
    f[u]=max(f[u],f[v]+1);
    g[u]=min(g[u],g[v]+1);
  }
  if(f[u]+g[u]<=mid) f[u]=-1e9;
  if(g[u]>mid&&b[u]) f[u]=max(f[u],0);
  if(f[u]==mid) f[u]=-1e9,g[u]=0,++tot;
}
int check(){
  tot=0;
  dfs(1,0);
  if(f[1]>=0) ++tot;
  return tot<=m;
}
int main(){
  n=read(),m=read();
  for(int i=1;i<=n;++i)b[i]=read();
  for(int i=1;i<n;++i){
    int x=read(),y=read();
    add(x,y);add(y,x);
  }
  int l=-1,r=n;
  while(l+1<r){
    mid=l+r>>1;
    check()?r=mid:l=mid;
  }
  printf("%d\n",r);
}