/ SeriousOJ /

Record Detail

Wrong Answer


  
# Status Time Cost Memory Cost
#1 Accepted 1ms 540.0 KiB
#2 Wrong Answer 1ms 540.0 KiB
#3 Wrong Answer 1ms 328.0 KiB

Code

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

void solve() {
    int n, k;
    cin >> n >> k;
    vector<int> v(n + 1);
    for(int i = 1; i <= n; i++) {
        ll x; cin >> x;
        v[i] = x;
    }
    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);
    }

    int root = 1;
    for(int i = 1; i <= n; i++) {
        if(g[i].size() == 1) {
            root = i;
            break;
        }
    }

    vector<vector<pair<ll, ll>>> dp(n + 1, vector<pair<ll, ll>> (k + 1, {-1, -1})); // {2nd max, 1st max}
    vector<bool> vis(n + 1);
    auto dfs = [&](auto&& self, int u) -> void {
        vis[u] = 1;
        dp[u][1].second = v[u];
        for(auto &child: g[u]) {
            if(vis[child]) continue;
            self(self, child);
            for(int i = 2; i <= k; i++) {
                ll mx = dp[child][i - 1].second;
                if(mx == -1) break;
                
                ll tmp = mx + v[u];
                auto &val = dp[u][i];
                if(val.second == -1) {
                    val.second = tmp;
                }
                else {
                    if(val.second < tmp) {
                        swap(val.first, val.second);
                        val.second = tmp;
                    }
                    else if(val.first < tmp) {
                        val.first = tmp;
                    }
                }
            }
        }
        return;
    };
    dfs(dfs, root);
    
    vis.assign(n + 1, 0);
    auto findAns = [&](auto&& self, int u) -> ll {
        vis[u] = 1;
        ll ans = 0;
        for(int i = 1; i <= k; i++) {
            int res = k - i;
            auto& val = dp[u][i];
            if(val.first == -1 && val.second == -1) continue;
            else if(val.first != -1 && val.second != -1) {
                ans = max(ans, val.first + val.second - v[u]);
            }
            else {
                ans = max(ans, val.second);
            }
        }
        
        for(auto &v: g[u]) {
            if(vis[v]) continue;
            self(self, v);
        }
        return ans;
    };
    cout << findAns(findAns, root) << endl;
    return;
}

int main() {
    ios::sync_with_stdio(false); cin.tie(0);
    int tc = 1;
    // cin >> tc;
    for (int t = 1; t <= tc; t++) {
        // cout << "Case " << t << ": ";
        solve();
    }
    return 0;
}

Information

Submit By
Type
Submission
Problem
P1161 Max path sum (Hard Version)
Language
C++17 (G++ 13.2.0)
Submit At
2025-02-20 17:32:37
Judged At
2025-02-20 17:32:37
Judged By
Score
2
Total Time
1ms
Peak Memory
540.0 KiB