BZOJ 2525: [Poi2011]Dynamite

bzoj_2525.cpp

#include<cstdio>
#include<vector>
#include<algorithm>
using namespace std;
const int kN = 300000+10;
int N, M;
vector<int> G[kN];
bool have_bomb[kN];
int T, need;
 
int Deepest[kN];
int Down[kN];
 
void DFS(int u, int fa) {
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i];
        if (v == fa) continue;
        DFS(v, u);
        if (Deepest[v] != -1) {
            Deepest[u] = max(Deepest[u], Deepest[v]+1);
        }
        if (Down[v] != -1) {
            Down[u] = max(Down[u], Down[v]-1);
        }
    }
     
    if (have_bomb[u]) Deepest[u] = max(Deepest[u], 0);
     
    if (Down[u] >= Deepest[u]) {
        Deepest[u] = -1;
    } else {
        if (Deepest[u] >= T) {
            Down[u] = T;
            Deepest[u] = -1;
            need++;
        }
    }
}
 
int CalcNeed(int time) {
    T = time;
    need = 0;
    for (int i = 1; i <= N; i++) Deepest[i] = Down[i] = -1;
    DFS(1, 0);
    if (Down[1] < Deepest[1]) need++;
    return need;
}
 
int main() {
    //freopen("in.txt", "r", stdin);
    scanf("%d %d", &N, &M);
    for (int i = 1; i <= N; i++) {
        int x; scanf("%d", &x);
        have_bomb[i] = x;
    }
    for (int i = 1; i <= N-1; i++) {
        int u, v; scanf("%d %d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
     
    int l = 0, r = N;
    while (l < r) {
        int mid = (l+r)/2;
        if (CalcNeed(mid) <= M) r = mid;
        else l = mid+1;
    }
 
    printf("%d\n", l);
    return 0;
}