8VC Venture Cup 2016 - Final Round (Div. 1 Edition) E - Preorder Test 树形dp

E - Preorder Test

思路：想到二分答案了之后就不难啦， 对于每个答案用树形dp取check, 如果二分的值是val, dp[ i ]表示 i 这棵子树答案不低于val的可以访问的

最多节点， 第二次dfs求出以每个点为根的答案。

#include<bits/stdc++.h>
#define LL long long
#define fi first
#define se second
#define mk make_pair
#define PII pair<int, int>
#define PLI pair<LL, int>
#define ull unsigned long long
using namespace std;

const int N = 2e5 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 1e9 + 7;
const double eps = 1e-8;

int n, k, a[N], dp[N], sz[N], mx[N], mx2[N];
vector<int> edge[N];

void dfs1(int u, int fa, int val) {
    sz[u] = 1;
    mx[u] = 0, mx2[u] = 0, dp[u] = 1;
    for(int v : edge[u]) {
        if(v == fa) continue;
        dfs1(v, u, val);
        sz[u] += sz[v];
        if(dp[v] == sz[v]) dp[u] += dp[v];
        else {
            if(dp[v] >= mx[u]) mx2[u] = mx[u], mx[u] = dp[v];
            else if(dp[v] > mx2[u]) mx2[u] = dp[v];
        }
    }
    dp[u] += mx[u];
    if(a[u] < val) dp[u] = 0;
}
void dfs2(int u, int fa, int cnt, int val, int &ans) {
    if(!dp[u]) {
        for(int v : edge[u]) {
            if(v == fa) continue;
            dfs2(v, u, 0, val, ans);
        }
    } else {
        int ret = dp[u];
        if(cnt == n - sz[u]) ret = max(ret, dp[u] + cnt);
        else if(cnt > mx[u]) ret = max(ret, dp[u] - mx[u] + cnt), mx2[u] = mx[u], mx[u] = cnt;
        else if(cnt > mx2[u]) mx2[u] = cnt;
        ans = max(ans, ret);
        for(int v : edge[u]) {
            if(v == fa) continue;
            if(dp[v] == sz[v]) dfs2(v, u, ret-dp[v], val, ans);
            else if(dp[v] == mx[u]) dfs2(v, u, ret-mx[u]+mx2[u], val, ans);
            else dfs2(v, u, ret, val, ans);
        }
    }
}

bool check(int val) {
    dfs1(1, 0, val);
    int ans = 0;
    dfs2(1, 0, 0, val, ans);
    return ans >= k;
}

int main() {
    scanf("%d%d", &n, &k);
    for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
    for(int i = 1; i < n; i++) {
        int u, v; scanf("%d%d", &u, &v);
        edge[u].push_back(v);
        edge[v].push_back(u);
    }
    int low = 1, high = 1000000, ans = 0;
    while(low <= high) {
        int mid = low + high >> 1;
        if(check(mid)) ans = mid, low = mid + 1;
        else high = mid - 1;
    }
    printf("%d\n", ans);
    return 0;
}

/*
*/