Submission #67966833

Source Code Expand

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll NEG = -(4LL << 60);         

int N, K;
vector<ll> A;
vector<vector<int>> g;

struct DP {
    vector<ll> closed;   
    vector<ll> open;    
};

DP dfs(int u, int p) {
    vector<ll> white(K + 1, NEG);
    vector<vector<ll>> open(2, vector<ll>(K + 1, NEG));
    vector<vector<ll>> closed(3, vector<ll>(K + 1, NEG));

    white[0]        = 0;           
    open[0][0]      = A[u];        
    closed[0][0]    = A[u];      

    for (int v : g[u]) if (v != p) {
        DP child = dfs(v, u);

        vector<ll> nwhite(K + 1, NEG);
        vector<vector<ll>> nopen(2, vector<ll>(K + 1, NEG));
        vector<vector<ll>> nclosed(3, vector<ll>(K + 1, NEG));

        for (int p1 = 0; p1 <= K; ++p1) if (white[p1] > NEG/2)
            for (int p2 = 0; p1 + p2 <= K; ++p2) if (child.closed[p2] > NEG/2)
                nwhite[p1 + p2] = max(nwhite[p1 + p2], white[p1] + child.closed[p2]);

        for (int k = 0; k <= 1; ++k)
            for (int p1 = 0; p1 <= K; ++p1) if (open[k][p1] > NEG/2)
                for (int p2 = 0; p1 + p2 <= K; ++p2) {
                    if (child.closed[p2] > NEG/2)                       
                        nopen[k][p1 + p2] = max(nopen[k][p1 + p2],
                                                open[k][p1] + child.closed[p2]);
                    if (child.open[p2] > NEG/2 && k + 1 <= 1)          
                        nopen[k + 1][p1 + p2] = max(nopen[k + 1][p1 + p2],
                                                     open[k][p1] + child.open[p2]);
                }

        for (int k = 0; k <= 2; ++k)
            for (int p1 = 0; p1 <= K; ++p1) if (closed[k][p1] > NEG/2)
                for (int p2 = 0; p1 + p2 <= K; ++p2) {
                    if (child.closed[p2] > NEG/2)                       
                        nclosed[k][p1 + p2] = max(nclosed[k][p1 + p2],
                                                   closed[k][p1] + child.closed[p2]);
                    if (child.open[p2] > NEG/2 && k + 1 <= 2)          
                        nclosed[k + 1][p1 + p2] = max(nclosed[k + 1][p1 + p2],
                                                        closed[k][p1] + child.open[p2]);
                }

        white.swap(nwhite);
        open.swap(nopen);
        closed.swap(nclosed);
    }

    DP res;
    res.closed.assign(K + 1, NEG);
    res.open.assign(K + 1, NEG);

    for (int p = 0; p <= K; ++p)
        res.open[p] = max(open[0][p], open[1][p]);      
    for (int p = 0; p <= K; ++p)
        res.closed[p] = max(res.closed[p], white[p]);

     for (int p = 0; p <= K; ++p) {
        if (p + 1 <= K && closed[0][p] > NEG/2)
            res.closed[p + 1] = max(res.closed[p + 1], closed[0][p]);

        if (p + 1 <= K && closed[1][p] > NEG/2)
            res.closed[p + 1] = max(res.closed[p + 1], closed[1][p]);

        if (p + 1 <= K && closed[2][p] > NEG/2)
            res.closed[p + 1] = max(res.closed[p + 1], closed[2][p]);
    }
    return res;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> N >> K;
    A.resize(N + 1);
    for (int i = 1; i <= N; ++i) cin >> A[i];

    g.assign(N + 1, {});
    for (int i = 1; i < N; ++i) {
        int u, v; cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    DP root = dfs(1, 0);
    ll ans = 0;
    for (int p = 0; p <= K; ++p) ans = max(ans, root.closed[p]);
    cout << ans << '\n';
    return 0;
}

Submission Info

Submission Time
Task F - Paint Tree 2
User Kifff12
Language C++ 20 (gcc 12.2)
Score 525
Code Size 3569 Byte
Status AC
Exec Time 289 ms
Memory 218848 KiB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 525 / 525
Status
AC × 3
AC × 25
Set Name Test Cases
Sample 00_sample_00.txt, 00_sample_01.txt, 00_sample_02.txt
All 00_sample_00.txt, 00_sample_01.txt, 00_sample_02.txt, 01_handmade_00.txt, 01_handmade_01.txt, 02_random_00.txt, 02_random_01.txt, 02_random_02.txt, 02_random_03.txt, 02_random_04.txt, 02_random_05.txt, 02_random_06.txt, 02_random_07.txt, 02_random_08.txt, 02_random_09.txt, 02_random_10.txt, 02_random_11.txt, 02_random_12.txt, 02_random_13.txt, 02_random_14.txt, 02_random_15.txt, 02_random_16.txt, 02_random_17.txt, 02_random_18.txt, 02_random_19.txt
Case Name Status Exec Time Memory
00_sample_00.txt AC 1 ms 3532 KiB
00_sample_01.txt AC 1 ms 3604 KiB
00_sample_02.txt AC 1 ms 3540 KiB
01_handmade_00.txt AC 1 ms 3456 KiB
01_handmade_01.txt AC 1 ms 3548 KiB
02_random_00.txt AC 22 ms 4916 KiB
02_random_01.txt AC 152 ms 13288 KiB
02_random_02.txt AC 38 ms 5976 KiB
02_random_03.txt AC 37 ms 6456 KiB
02_random_04.txt AC 154 ms 13788 KiB
02_random_05.txt AC 192 ms 15924 KiB
02_random_06.txt AC 191 ms 15952 KiB
02_random_07.txt AC 196 ms 16068 KiB
02_random_08.txt AC 199 ms 16000 KiB
02_random_09.txt AC 194 ms 16076 KiB
02_random_10.txt AC 175 ms 16532 KiB
02_random_11.txt AC 174 ms 16524 KiB
02_random_12.txt AC 254 ms 181416 KiB
02_random_13.txt AC 289 ms 218848 KiB
02_random_14.txt AC 268 ms 154608 KiB
02_random_15.txt AC 285 ms 170184 KiB
02_random_16.txt AC 273 ms 171240 KiB
02_random_17.txt AC 286 ms 171244 KiB
02_random_18.txt AC 178 ms 15736 KiB
02_random_19.txt AC 195 ms 15696 KiB