Subtask ~2~: ~n \le 400~

Gọi ~dp_{u,h}~ là vẻ đẹp lớn nhất khi xét tới cây gốc ~u~ và đỉnh được chọn gần ~u~ nhất cách ~u~ một khoảng ~h~, chúng ta sẽ gộp các con của ~v~ để tính.

Code tham khảo

#include <bits/stdc++.h>
using namespace std;
//(>'-'<)
typedef long long ll;

const int N = 5e3 + 2;
const int oo = 1e9 + 7;

int n, k, a[N];
int sz[N], dp[N][N], tmp[N];
vector<int> adj[N];

bool maximize(int &a, int b) { return a < b ? a = b, true : false; }

void dfs(int u, int p) {
    sz[u] = 1;
    dp[u][0] = a[u];
    for(int &v : adj[u]) if(v != p) {
        dfs(v, u);
        for(int h = 0; h <= sz[u] + sz[v]; ++h) {
            if (h > 0) tmp[h] = max(dp[u][h], dp[v][h - 1]);
            else tmp[h] = dp[u][h];
        }
        for(int h1 = 0; h1 <= sz[u]; ++h1) {
            for(int h2 = k - h1; h2 <= sz[v]; ++h2) {
                maximize(tmp[min(h1, h2 + 1)], dp[u][h1] + dp[v][h2]);
            }
        }
        sz[u] += sz[v];
        for(int h = 0; h <= sz[u]; ++h) {
            dp[u][h] = tmp[h];
            tmp[h] = -oo;
        }
    }
}

int main() {
    ios_base::sync_with_stdio(0);cin.tie(0);

    cin >> n >> k;
    for(int i = 0; i < n; ++i) cin >> a[i];
    for(int i = 1; i < n; ++i) {
        int x; cin >> x;
        adj[i].push_back(x);
        adj[x].push_back(i);
    }
    memset(dp, -0x3f, sizeof dp);
    memset(tmp, -0x3f, sizeof tmp);
    dfs(0, -1);
    int res = 0;
    for(int h = 0; h <= n; ++h) {
        res = max(res, dp[0][h]);
    }
    cout << res;

    return 0;
}

Độ phức tạp: ~\mathcal{O}({n^3})~.

Subtask ~3~: ~n \le 5000~

Ta có thể giảm độ phức tạp khi tính ~dp_{u,h}~ bằng cách sử dụng suffix max.

Code tham khảo

int n, k, a[N];
int sz[N], dp[N][N], tmp[N], suf_max_u[N], suf_max_v[N];
vector<int> adj[N];

bool maximize(int &a, int b) { return a < b ? a = b, true : false; }

void dfs(int u, int p) {
    sz[u] = 1;
    dp[u][0] = a[u];
    for(int &v : adj[u]) if(v != p) {
        dfs(v, u);
        for(int h = 0; h <= sz[u] + sz[v]; ++h) {
            if (h > 0) tmp[h] = max(dp[u][h], dp[v][h - 1]);
            else tmp[h] = dp[u][h];
        }
        suf_max_u[sz[u] + 1] = -oo;
        for(int h = sz[u]; h >= 0; --h) {
            suf_max_u[h] = max(dp[u][h], suf_max_u[h + 1]);
        }
        suf_max_v[sz[v] + 1] = -oo;
        for(int h = sz[v]; h >= 0; --h) {
            suf_max_v[h] = max(dp[v][h], suf_max_v[h + 1]);
        }
        for(int h1 = 0; h1 <= sz[u]; ++h1) {
            int h2 = max(k - h1, h1 - 1);
            if (h2 <= sz[v]) maximize(tmp[h1], dp[u][h1] + suf_max_v[h2]);
        }
        for(int h2 = 0; h2 <= sz[v]; ++h2) {
            int h1 = max(k - h2, h2 + 1);
            if (h1 <= sz[u]) maximize(tmp[h2 + 1], suf_max_u[h1] + dp[v][h2]);
        }
        sz[u] += sz[v];
        for(int h = 0; h <= sz[u]; ++h) {
            dp[u][h] = tmp[h];
            tmp[h] = -oo;
        }
    }
}

Subtask ~4~: ~n \le 5\times 10^5~ và mỗi đỉnh có không quá ~2~ cạnh nối.

Gọi ~dp_i~ là độ đẹp lớn nhất khi xét tới đỉnh ~i~, ~dp_i = max(dp_{i-1}, dp_{i - k - 1} + a_{i})~

Subtask ~5~: ~n \le 5\times 10^5~.

Để ý khi gộp các đỉnh con của ~u~, độ phức tạp sẽ là ~\mathcal{O}({size_u + size_v})~. Chúng ta có thể thực hiện việc gộp trong ~\mathcal{O}({min(size_u, size_v)})~:

• Nếu ~size_u \leq size_v~, chúng ta thực hiện gộp ~dp_u~ vào ~dp_v~. Ngược lại, gộp ~dp_v~ vào ~dp_u~.
• Chỉ chạy for tới ~min(size_u, size_v)~ vì các trạng thái lớn hơn không bị ảnh hưởng.

Mỗi trạng thái ~h~ của ~dp~ chỉ lặp lại nhiều nhất ~\mathcal{O}({logn})~ khi tính, việc cài đặt khéo léo sẽ giúp chúng ta giảm độ phức tạp thành ~\mathcal{O}({nlogn})~.

Ngoài ra có thể giảm độ phức tạp còn ~\mathcal{O}({n})~, tham khảo tại: https://codeforces.com/blog/entry/70822