Submission #1310513

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

Main

#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 from Root ---
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 ---
// Path(u) = u -> Rev(v1) -> Rev(v2) ...
void print_k3(int u, int p, bool reverse_mode) {
    if (!reverse_mode) {
        // Forward: u -> Rev(v1) -> Rev(v2) ...
        cout << u << " ";
        for (int v : adj[u]) {
            if (v != p) {
                print_k3(v, u, true);
            }
        }
    } else {
        // Reverse: ... v2 -> v1 -> u
        // Rev(Path(u)) = ... Path(v2) -> Path(v1) -> u
        // We iterate children in reverse order of how they were added in forward pass
        // Since order doesn't strictly matter for profit, we just reverse the iteration
        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 Solution: Tree DP ---
struct DPState {
    long long val;
    int choice_a, choice_b, choice_c; // Indices of children used for d2, d2, d1
};

// dp[u][0] = dp1 (start u)
// dp[u][1] = dp2 (start u, end shallow)
// dp[u][2] = dp3 (start child)
DPState dp[200005][3];

// Helper to calculate gain
struct Cand {
    long long gain;
    int id; // Child node index
};

bool compareCand(const Cand& a, const Cand& b) {
    return a.gain > b.gain;
}

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];

    // Collect candidates
    vector<Cand> cands_d1, cands_d2;
    for (int v : children) {
        cands_d1.push_back({dp[v][0].val - P[v], v});
        cands_d2.push_back({dp[v][1].val - P[v], v});
    }

    sort(cands_d1.begin(), cands_d1.end(), compareCand);
    sort(cands_d2.begin(), cands_d2.end(), compareCand);

    // Limit to top 3 (sufficient for finding disjoint triplet)
    if (cands_d1.size() > 3) cands_d1.resize(3);
    if (cands_d2.size() > 3) cands_d2.resize(3);

    // --- DP1: Start u ---
    // Max: base + D2[a] + D1[b] (a != b)
    dp[u][0] = {base, -1, -1, -1}; // Default: no special children
    
    // Case 0: No specials (already set)
    // Case 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};
    }
    // Case 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};
    }
    // Case 3: Both
    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};
                }
            }
        }
    }

    // --- DP2: Start u, End shallow ---
    // Max: base + D2[a]
    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 ---
    // Max: base + D2[a] + D2[b] + D1[c] (distinct)
    dp[u][2] = {-INF, -1, -1, -1}; // Valid dp3 MUST start at a child
    
    // Iterate combinations. a = start child (D2), b = rev child (D2), c = end child (D1)
    // Note: a is the 'start' child. It MUST exist.
    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
        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
        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};
                    }
                }
            }
        }
    }
}

// Reconstruction for K=2
vector<int> path_k2;
void reconstruct_k2(int u, int type) {
    DPState& s = dp[u][type];
    
    // Identify special children
    int va = s.choice_a; // The 'start' or 'loop' D2 child
    int vb = s.choice_b; // The second 'loop' D2 child (only for dp3)
    int vc = s.choice_c; // The D1 child
    
    // Order of operations:
    // If DP1: u -> Rev(D2[va]) -> singles -> D1[vc]
    // If DP2: u -> Rev(D2[va]) -> singles
    // If DP3: D2[va] -> u -> Rev(D2[vb]) -> singles -> D1[vc]
    
    auto do_singles = [&](int exclude1, int exclude2, int exclude3) {
        for(int v : adj[u]) {
            // Check if v is parent (hacky, using P[v] check works if tree rooted at 1?)
            // We need to pass parent or check adjacency.
            // Since we stored children in DFS, we iterate original adj and skip specials.
            bool is_child = true; 
            // In reconstruction we don't have 'p'. But flow is directed.
            // Just need to ensure we don't go UP.
            // However, special children are definitely children.
            // 'singles' should be children only.
            // Just checking against specials is enough if we only traverse children.
        }
        // Actually simpler: iterate adj[u], if v not parent and not special, print v.
    };

    // Need access to children list to print singles.
    // We can assume 'adj' is undirected, we need 'p'.
    // Or just check against parent in global.
}

// Recursive print with parent tracking
void print_k2_rec(int u, int p, int type) {
    DPState& s = dp[u][type];
    int va = s.choice_a; 
    int vb = s.choice_b;
    int vc = s.choice_c;

    auto print_singles = [&](int except_a, int except_b, int except_c) {
        for (int v : adj[u]) {
            if (v != p && v != except_a && v != except_b && v != except_c) {
                cout << v << " ";
            }
        }
    };

    if (type == 0) { // DP1
        cout << u << " ";
        if (va != -1) {
             // Rev(D2) means: reconstruct D2 then reverse? 
             // D2 path: u -> ... -> child.
             // Rev(D2): child -> ... -> u.
             // But here we are at u. We need to go u -> Rev(D2 of child).
             // D2 of child starts at child.
             // So Rev(D2 of child) starts at leaf of child, ends at child.
             // Wait. Logic check.
             // Transition: u -> Rev(dp2[v]).
             // Means u -> leaf -> ... -> v.
             // This works.
             // So we need a function to print D2 in reverse.
             // print_rev_d2(va)
        }
        // Actually, let's use a vector to collect path then reverse/print.
    }
}

// We'll use a vector builder for simplicity in logic
void build_k2(int u, int p, int type, vector<int>& path) {
    DPState& s = dp[u][type];
    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);
            }
        }
    };

    // Helper to add Rev(DP2[v])
    auto add_rev_dp2 = [&](int v) {
        vector<int> sub;
        build_k2(v, u, 1, sub);
        // Reverse sub and append
        path.insert(path.end(), sub.rbegin(), sub.rend());
    };

    // Helper to add DP2[v] (Forward)
    auto add_dp2 = [&](int v) {
        build_k2(v, u, 1, path);
    };

    // Helper to add DP1[v]
    auto add_dp1 = [&](int v) {
        build_k2(v, u, 0, path);
    };

    if (type == 0) { // DP1: u -> Rev(va) -> singles -> vc
        path.push_back(u);
        if (va != -1) add_rev_dp2(va);
        add_singles();
        if (vc != -1) add_dp1(vc);
    } 
    else if (type == 1) { // DP2: u -> Rev(va) -> singles
        path.push_back(u);
        if (va != -1) add_rev_dp2(va);
        add_singles();
    } 
    else if (type == 2) { // DP3: va -> u -> Rev(vb) -> singles -> vc
        if (va != -1) add_dp2(va);
        path.push_back(u);
        if (vb != -1) add_rev_dp2(vb);
        add_singles();
        if (vc != -1) add_dp1(vc);
    }
}


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