Submission #1310515

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

Main

/*
 * Problem: Travelling Trader (CCO 2023 P2)
 * Solution by: Gemini
 * Logic:
 * - K=1: Longest path from root (DFS).
 * - K=3: Euler Tour (Visit all nodes).
 * - K=2: Tree DP with State [Start][End_Type].
 * - dp[u][0] (DP1): Start u, stay in subtree.
 * - dp[u][1] (DP2): Start u, end at u or immediate child.
 * - dp[u][2] (DP3): Start at a child of u, stay in subtree.
 * - Optimization: Use Top 3 candidates to avoid O(N^2) on star graphs.
 */

#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 Solution =================
// Longest path starting at 1
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 Solution =================
// Visit all nodes via Euler Tour logic
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); // Reverse the child's tour
            }
        }
    } else {
        vector<int> children;
        for (int v : adj[u]) if (v != p) children.push_back(v);
        // Visit children in reverse order to maintain flow
        for (int i = children.size() - 1; i >= 0; i--) {
            print_k3(children[i], u, false); // Keep child's internal tour forward? 
            // Actually standard strategy: u -> Rev(v1) -> Rev(v2).
            // When reversing u, we do ... v2 -> v1 -> u.
        }
        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 Solution =================
struct DPState {
    long long val;
    int choice_a, choice_b, choice_c; 
    // a, b: D2 children (Loops), c: D1 child (Tail)
    // Special: choice_a = -2 indicates jump to Grandchild (P[u] + dp3[v])
};

DPState dp[200005][3];

struct Cand {
    long long gain;
    int id;
};

// Optimization: Maintain only Top 3 candidates
void update_top3(vector<Cand>& top3, const 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;
    long long max_dp3_jump = -INF;
    int dp3_jump_id = -1;

    for (int v : children) {
        update_top3(cands_d1, {dp[v][0].val - P[v], v}); // Gain if used as D1
        update_top3(cands_d2, {dp[v][1].val - P[v], v}); // Gain if used as D2
        
        // Track jump to grandchild: P[u] + dp3[v]
        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};
    
    // Combinations of D2(a) and D1(b)
    // 1. Just D1(b)
    for(auto& b : cands_d1) {
        if(base + b.gain > dp[u][0].val) dp[u][0] = {base + b.gain, -1, -1, b.id};
    }
    // 2. Just D2(a)
    for(auto& a : cands_d2) {
        if(base + a.gain > dp[u][0].val) dp[u][0] = {base + a.gain, a.id, -1, -1};
    }
    // 3. Both D2(a) + D1(b)
    for(auto& a : cands_d2) {
        for(auto& b : cands_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};
                }
            }
        }
    }
    // 4. Jump to Grandchild (Overrides singles)
    if (max_dp3_jump > dp[u][0].val) {
        dp[u][0] = {max_dp3_jump, -2, -1, dp3_jump_id};
    }

    // --- DP2: Start u, End shallow ---
    dp[u][1] = {base, -1, -1, -1};
    for(auto& a : cands_d2) {
        if(base + a.gain > dp[u][1].val) dp[u][1] = {base + a.gain, a.id, -1, -1};
    }

    // --- DP3: Start child ---
    dp[u][2] = {-INF, -1, -1, -1};
    
    for(auto& a : cands_d2) {
        long long current = base + a.gain;
        // Just a
        if(current > dp[u][2].val) dp[u][2] = {current, a.id, -1, -1};

        // a + b (Two D2 loops)
        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 + c (D2 start, D1 end)
        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 + b + c
        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};
                    }
                }
            }
        }
    }
}

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];
    
    // Case: Jump to Grandchild
    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) -> singles -> vc
        path.push_back(u);
        if (va != -1) add_child_path(va, u, 1, path, true); // Rev D2
        add_singles();
        if (vc != -1) add_child_path(vc, u, 0, path, false); // D1
    } 
    else if (type == 1) { // DP2: u -> Rev(va) -> singles
        path.push_back(u);
        if (va != -1) add_child_path(va, u, 1, path, true); // Rev D2
        add_singles();
    } 
    else if (type == 2) { // DP3: va -> u -> Rev(vb) -> singles -> vc
        if (va != -1) add_child_path(va, u, 1, path, false); // Forward D2 (Start child)
        path.push_back(u);
        if (vb != -1) add_child_path(vb, u, 1, path, true); // Rev D2
        add_singles();
        if (vc != -1) add_child_path(vc, u, 0, path, false); // 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