Submission #1310516

#	Time	Username	Problem	Language	Result	Execution time	Memory
1310516		ayuxhkumxr22	Travelling Trader (CCO23_day2problem2)	C++20	9 / 25	2091 ms	44868 KiB

Main

/*
 * Problem: Travelling Trader (CCO 2023 P2)
 * Solution: Optimized Tree DP for K=1, 2, 3
 * Fixes for K=2:
 * - Proper handling of Incoming candidates (dp2 vs dp3).
 * - Strict definition of Loop candidates (dp2 only).
 * - Inclusion of Grandchild Jump for dp1.
 * - O(N) transitions using Top-3.
 */

#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>

using namespace std;

const long long INF = 1e18;
int N, K;
vector<long long> P;
vector<vector<int>> adj;

// ================= K=1: Longest Path =================
long long max_profit_k1 = -1;
vector<int> path_k1;

void dfs_k1(int u, int p, long long current_profit, vector<int>& current_path) {
    current_path.push_back(u);
    current_profit += P[u];
    
    if (current_profit > max_profit_k1) {
        max_profit_k1 = current_profit;
        path_k1 = current_path;
    }
    
    for (int v : adj[u]) {
        if (v != p) {
            dfs_k1(v, u, current_profit, current_path);
        }
    }
    current_path.pop_back();
}

void solve_k1() {
    vector<int> current_path;
    dfs_k1(1, 0, 0, current_path);
    cout << max_profit_k1 << "\n";
    cout << path_k1.size() << "\n";
    for (int i = 0; i < path_k1.size(); i++) cout << path_k1[i] << (i == path_k1.size()-1 ? "" : " ");
    cout << "\n";
}

// ================= K=3: Euler Tour =================
void print_k3(int u, int p, bool reverse_mode) {
    if (!reverse_mode) {
        cout << u << " ";
        for (int v : adj[u]) {
            if (v != p) {
                print_k3(v, u, true);
            }
        }
    } else {
        vector<int> children;
        for (int v : adj[u]) if (v != p) children.push_back(v);
        for (int i = children.size() - 1; i >= 0; i--) {
            print_k3(children[i], u, false);
        }
        cout << u << " ";
    }
}

void solve_k3() {
    long long total = 0;
    for (int i = 1; i <= N; i++) total += P[i];
    cout << total << "\n";
    cout << N << "\n";
    print_k3(1, 0, false);
    cout << "\n";
}

// ================= K=2: Complex Tree DP =================

struct DPState {
    long long val;
    int choice_a, choice_b, choice_c; 
    // a: Incoming/Loop, b: Loop, c: Tail
    // Special: choice_a = -2 -> Grandchild Jump
    int type_a; // 1 for dp2, 2 for dp3 (Only for choice_a)
};

DPState dp[200005][3];

struct Cand {
    long long gain;
    int id;
    int type; // 1 for dp2, 2 for dp3
};

void update_top3(vector<Cand>& top3, Cand c) {
    top3.push_back(c);
    int i = top3.size() - 1;
    while(i > 0) {
        if(top3[i].gain > top3[i-1].gain) {
            swap(top3[i], top3[i-1]);
            i--;
        } else break;
    }
    if(top3.size() > 3) top3.pop_back();
}

void dfs_k2(int u, int p) {
    vector<int> children;
    long long sum_p = 0;
    
    for (int v : adj[u]) {
        if (v != p) {
            dfs_k2(v, u);
            children.push_back(v);
            sum_p += P[v];
        }
    }

    long long base = sum_p + P[u];

    vector<Cand> cands_d1, cands_d2, cands_in;
    long long max_dp3_jump = -INF;
    int dp3_jump_id = -1;

    for (int v : children) {
        // Candidates for Tail (D1)
        update_top3(cands_d1, {dp[v][0].val - P[v], v, 0});
        
        // Candidates for Loop (Strict D2)
        update_top3(cands_d2, {dp[v][1].val - P[v], v, 1});
        
        // Candidates for Incoming (Max of Strict D2, D3)
        // Incoming allows ending at Child (D3) or U (D2)
        long long gain_d2 = dp[v][1].val - P[v];
        long long gain_d3 = dp[v][2].val - P[v];
        
        if (gain_d3 > gain_d2) {
            update_top3(cands_in, {gain_d3, v, 2});
        } else {
            update_top3(cands_in, {gain_d2, v, 1});
        }

        // Grandchild Jump (Standalone path)
        if (P[u] + dp[v][2].val > max_dp3_jump) {
            max_dp3_jump = P[u] + dp[v][2].val;
            dp3_jump_id = v;
        }
    }

    // --- DP1: Start u ---
    dp[u][0] = {base, -1, -1, -1, 0};
    
    // 1. Loop + Tail
    for(auto& a : cands_d2) { // Loop must be Strict D2
        for(auto& b : cands_d1) { // Tail is D1
            if(a.id != b.id) {
                if(base + a.gain + b.gain > dp[u][0].val) {
                    dp[u][0] = {base + a.gain + b.gain, a.id, -1, b.id, 1};
                }
            }
        }
    }
    // 2. Just Loop (D2)
    for(auto& a : cands_d2) {
        if(base + a.gain > dp[u][0].val) dp[u][0] = {base + a.gain, a.id, -1, -1, 1};
    }
    // 3. Just Tail (D1)
    for(auto& b : cands_d1) {
        if(base + b.gain > dp[u][0].val) dp[u][0] = {base + b.gain, -1, -1, b.id, 0};
    }
    
    // 4. Grandchild Jump (Overrides all)
    if (max_dp3_jump > dp[u][0].val) {
        dp[u][0] = {max_dp3_jump, -2, -1, dp3_jump_id, 0};
    }

    // --- DP2: Start u, End u (Strict) ---
    dp[u][1] = {base, -1, -1, -1, 0};
    // Only Loop allowed (Tail would end deep)
    for(auto& a : cands_d2) {
        if(base + a.gain > dp[u][1].val) dp[u][1] = {base + a.gain, a.id, -1, -1, 1};
    }

    // --- DP3: Start Child ---
    dp[u][2] = {-INF, -1, -1, -1, 0};
    
    for(auto& a : cands_in) { // Incoming (D2 or D3)
        long long current = base + a.gain;
        
        // Just Incoming
        if(current > dp[u][2].val) dp[u][2] = {current, a.id, -1, -1, a.type};

        // Incoming + Loop
        for(auto& b : cands_d2) {
            if(a.id != b.id) {
                if(current + b.gain > dp[u][2].val) {
                    dp[u][2] = {current + b.gain, a.id, b.id, -1, a.type};
                }
            }
        }
        // Incoming + Tail
        for(auto& c : cands_d1) {
            if(a.id != c.id) {
                if(current + c.gain > dp[u][2].val) {
                    dp[u][2] = {current + c.gain, a.id, -1, c.id, a.type};
                }
            }
        }
        // Incoming + Loop + Tail
        for(auto& b : cands_d2) {
            if(a.id == b.id) continue;
            for(auto& c : cands_d1) {
                if(a.id != c.id && b.id != c.id) {
                    if(current + b.gain + c.gain > dp[u][2].val) {
                        dp[u][2] = {current + b.gain + c.gain, a.id, b.id, c.id, a.type};
                    }
                }
            }
        }
    }
}

// Reconstruction
void build_k2(int u, int p, int type, vector<int>& path);

void add_child_path(int v, int u, int type, vector<int>& path, bool reverse_output) {
    vector<int> sub;
    build_k2(v, u, type, sub);
    if(reverse_output) {
        path.insert(path.end(), sub.rbegin(), sub.rend());
    } else {
        path.insert(path.end(), sub.begin(), sub.end());
    }
}

void build_k2(int u, int p, int type, vector<int>& path) {
    DPState& s = dp[u][type];
    
    // Grandchild Jump
    if (type == 0 && s.choice_a == -2) {
        path.push_back(u);
        add_child_path(s.choice_c, u, 2, path, false); // Forward DP3
        return;
    }

    int va = s.choice_a;
    int vb = s.choice_b;
    int vc = s.choice_c;

    auto add_singles = [&]() {
        for (int v : adj[u]) {
            if (v != p && v != va && v != vb && v != vc) {
                path.push_back(v);
            }
        }
    };

    if (type == 0) { // DP1: u -> Rev(va, D2) -> Singles -> vc(D1)
        path.push_back(u);
        if (va != -1) add_child_path(va, u, 1, path, true);
        add_singles();
        if (vc != -1) add_child_path(vc, u, 0, path, false);
    } 
    else if (type == 1) { // DP2: u -> Rev(va, D2) -> Singles
        path.push_back(u);
        if (va != -1) add_child_path(va, u, 1, path, true);
        add_singles();
    } 
    else if (type == 2) { // DP3: va(In) -> u -> Rev(vb, D2) -> Singles -> vc(D1)
        if (va != -1) {
            // Incoming can be D2 or D3
            if (s.type_a == 1) add_child_path(va, u, 1, path, false); // D2 Forward
            else add_child_path(va, u, 2, path, true); // Rev(Rev(D3)) = D3 Reversed? 
            // Wait. Incoming is Rev(Rev(path))? 
            // Incoming structure: Path -> u.
            // D2(Forward) is u->child. Rev(D2) is child->u. This works.
            // D3(Forward) is child->child. Rev(D3) is child->child.
            // But we need to END at u.
            // Dist(End(Rev(D3)), u) = Dist(child, u) = 2. Valid.
            // So we need Rev(D3).
            // D3 standard output is Forward. So we need Reverse output.
        }
        
        path.push_back(u);
        
        if (vb != -1) add_child_path(vb, u, 1, path, true); // Loop Rev(D2)
        add_singles();
        if (vc != -1) add_child_path(vc, u, 0, path, false); // Tail D1
    }
}

void solve_k2() {
    dfs_k2(1, 0);
    cout << dp[1][0].val << "\n";
    vector<int> path;
    build_k2(1, 0, 0, path);
    cout << path.size() << "\n";
    for(size_t i=0; i<path.size(); ++i) cout << path[i] << (i==path.size()-1?"":" ");
    cout << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    if (!(cin >> N >> K)) return 0;
    P.resize(N + 1);
    adj.resize(N + 1);
    for (int i = 0; i < N - 1; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for (int i = 1; i <= N; i++) cin >> P[i];

    if (K == 1) solve_k1();
    else if (K == 2) solve_k2();
    else solve_k3();

    return 0;
}

• Subtask 1: 0 / 2
• Subtask 2: 0 / 7
• Subtask 3: 0 / 3
• Subtask 4: 0 / 4
• Subtask 5: 4 / 4
• Subtask 6: 5 / 5