题目

Luogu
LOJ
Acwing

思路

代码

#include <iostream>
#include <algorithm>
#include <cstring>
using namespace std;
const int N = 500010, M = 22;
int n, m, d, w[N], st[N];
int h[N], val[N << 1], ptr[N << 1], idx;
int f[N][M], g[N][M];
void add(int a, int b) { val[idx] = b, ptr[idx] = h[a], h[a] = idx++; }
void DFS(int u, int fa) {
    if (st[u]) f[u][0] = g[u][0] = w[u];
    else f[u][0] = g[u][0] = 0;
    for (int i = 1; i <= d; i++) f[u][i] = w[u];
    f[u][d + 1] = 0x3f3f3f3f;
    for (int i = h[u], v = val[i]; i != -1; i = ptr[i], v = val[i]) {
        if (v == fa) continue;
        DFS(v, u);
        // 这里当 j == d 的时候, f[u][j + 1] 即 f[u][d + 1] = INF
        for (int j = d; j >= 0; j--)
            f[u][j] = min(f[u][j] + g[v][j], g[u][j + 1] + f[v][j + 1]);
        for (int j = d; j >= 0; j--) f[u][j] = min(f[u][j], f[u][j + 1]);
        g[u][0] = f[u][0];
        for (int j = 1; j <= d + 1; j++) g[u][j] += g[v][j - 1];
        for (int j = 1; j <= d + 1; j++) g[u][j] = min(g[u][j], g[u][j - 1]);
    }
}
int main() {
    cin >> n >> d;
    memset(h, -1, sizeof h);
    for (int i = 1; i <= n; i++) cin >> w[i];
    cin >> m;
    for (int i = 1, t; i <= m; i++) 
        cin >> t, st[t] = true;
    for (int i = 1, a, b; i < n; i++)
         cin >> a >> b, add(a, b), add(b, a);
    DFS(1, 0);
    cout << f[1][0] << endl;
    return 0;
}