1. 程式人生 > 實用技巧 >牛客 msc的寵物 【二分+樹形dp】

牛客 msc的寵物 【二分+樹形dp】

題目連結
自己假自己,怎麼寫都只能過一半的測試點,學了別人的題解。
題目描述
給定一棵樹,每個結點有一個權值,求\(k+1\)個連通塊中最值之差的最小值是多少。
思路
二分答案很顯然,主要是用樹形dp計算最少需要的次數。
\(dp[u][i]\):結點\(u\)的值\(a[u]\)在範圍為\([a_i,a_i + mid]\)中,且滿足其子樹所有點的值在該範圍中時需要刪除的邊數。
\(f[u]\):結點\(u\)所需要刪除的最少的邊的數量。
轉移方程:\(dp[u][i]=min(dp[v][i],f[v]+1)\)
程式碼

#include<bits/stdc++.h>
using namespace std;

typedef long long LL;
const int inf = 0x3f3f3f3f;
const int N = 1010;
LL dp[N][N], f[N];
LL a[N];
vector<int> e[N];
int n, k;

void dfs(int u, int fa, LL mid) {
    for(int i = 1; i <= n; i++) {
        if(a[i] <= a[u] && a[u] <= a[i] + mid) {
            dp[u][i] = 0;
        } else dp[u][i] = inf;
    }
    for(int i = 0; i < e[u].size(); i++) {
        int v = e[u][i];
        if(v == fa) continue;
        dfs(v, u, mid);
        for(int j = 1; j <= n; j++) {
            dp[u][j] += min(dp[v][j], f[v] + 1);
        }
    }
    for(int i = 1; i <= n; i++) {
        f[u] = min(f[u], dp[u][i]);
    }
}

bool check(LL mid) {
    memset(f, 0x3f, sizeof f);
    dfs(1, 0, mid);
    return f[1] <= k;
}

void solve() {
    scanf("%d%d", &n, &k);
    for(int i = 1; i <= n; i++) {
        scanf("%lld", &a[i]);
    }
    for(int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    LL l = 0, r = 2e9;
    LL res = -1;
    while(l <= r) {
        LL mid = l + r >> 1;
//        printf("mid = %lld\n", mid);
        if(check(mid)) {
            res = mid;
            r = mid - 1;
        } else l = mid + 1;
    }
    printf("%lld\n", res);
}

int main() {
//    freopen("in.txt", "r", stdin);
    solve();
    return 0;
}