/ SeriousOJ /

Record Detail

Wrong Answer


  
# Status Time Cost Memory Cost
#1 Accepted 1ms 540.0 KiB
#2 Accepted 1ms 540.0 KiB
#3 Accepted 1ms 540.0 KiB
#4 Wrong Answer 1ms 540.0 KiB
#5 Accepted 1ms 540.0 KiB
#6 Accepted 1ms 540.0 KiB
#7 Accepted 1ms 540.0 KiB
#8 Wrong Answer 1ms 540.0 KiB

Code

#include<bits/stdc++.h>
#define endl '\n'
using namespace std;
using ll = long long;

int main() {
    ios::sync_with_stdio(false); cin.tie(0);
    int n, k;
    cin >> n >> k;
    vector<int> a(n + 1);
    for(int i = 1; i <= n; i++) cin >> a[i];
    vector<vector<int>> g(n + 1);
    for(int i = 0; i + 1 < n; i++) {
        int u, v; cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    vector<vector<priority_queue<pair<ll, int>, vector<pair<ll, int>>, greater<>>>> dp(n + 1, vector<priority_queue<pair<ll, int>, vector<pair<ll, int>>, greater<>>> (k + 1)); 
    auto dfs = [&](auto&& self, int u, int par) -> ll {
        dp[u][1].push({a[u], -1});
        int dir = 0;
        ll mx = 0;
        for(auto &v: g[u]) {
            if(v == par) continue;
            mx = max(mx, self(self, v, u));
            for(int i = 2; i <= k; i++) {
                if(dp[v][i - 1].size() == 0) break;
                if(dp[v][i - 1].size() == 1) dp[u][i].push({a[u] + dp[v][i - 1].top().first, dir});
                else if(dp[v][i - 1].size() == 2) {
                    auto x = dp[v][i - 1].top();
                    dp[v][i - 1].pop();
                    dp[u][i].push({a[u] + dp[v][i - 1].top().first, dir});
                    dp[v][i - 1].push(x);
                }
                if(dp[u][i].size() > 2) dp[u][i].pop();
            }
            ++dir;
        }
        // cout << "==== " << u << " ===\n";
        // for(int i = 1; i <= k; i++) {
        //     cout << i << "=> ";
        //     auto x = dp[u][i];
        //     while(x.size()) {
        //         cout << "{" << x.top().first << " " << x.top().second << "}, ";
        //         x.pop();
        //     }
        //     cout << endl;
        // }
        for(int i = 1; i <= k; i++) {
            if(dp[u][i].size() && dp[u][k - i + 1].size()) {
                if(dp[u][i].top().second != dp[u][k - i + 1].top().second) {
                    mx = max(mx, dp[u][i].top().first + dp[u][k - i + 1].top().first - a[u]);
                }
                if(dp[u][i].size() == 2) {
                    auto tmp = dp[u][i].top();
                    dp[u][i].pop();
                    if(dp[u][i].top().second != dp[u][k - i + 1].top().second) {
                        mx = max(mx, dp[u][i].top().first + dp[u][k - i + 1].top().first - a[u]);
                    }
                    dp[u][i].push(tmp);
                }
                if(dp[u][k - i + 1].size() == 2) {
                    auto tmp = dp[u][k - i + 1].top();
                    dp[u][k - i + 1].pop();
                    if(dp[u][i].top().second != dp[u][k - i + 1].top().second) {
                        mx = max(mx, dp[u][i].top().first + dp[u][k - i + 1].top().first - a[u]);
                    }
                    dp[u][k - i + 1].push(tmp);
                }
                if(dp[u][i].size() == 2 && dp[u][k - i + 1].size() == 2) {
                    auto tmp1 = dp[u][i].top();
                    dp[u][i].pop();
                    auto tmp2 = dp[u][k - i + 1].top();
                    dp[u][k - i + 1].pop();
                    if(dp[u][i].top().second != dp[u][k - i + 1].top().second) {
                        mx = max(mx, dp[u][i].top().first + dp[u][k - i + 1].top().first - a[u]);
                    }
                    dp[u][i].push(tmp1);
                    dp[u][k - i + 1].push(tmp2);
                }
            }
        }
        // cout << u << "=> " << mx << endl;
        return mx;
    };
    cout << dfs(dfs, 1, 0) << endl;
    return 0;
}

Information

Submit By
Type
Submission
Problem
P1161 Max path sum (Hard Version)
Language
C++17 (G++ 13.2.0)
Submit At
2025-03-24 19:09:34
Judged At
2025-03-24 19:09:34
Judged By
Score
12
Total Time
1ms
Peak Memory
540.0 KiB