Submission #1310518

#	Time	Username	Problem	Language	Result	Execution time	Memory
1310518		ayuxhkumxr22	Travelling Trader (CCO23_day2problem2)	C++20	9 / 25	2090 ms	44880 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 =================
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 =================
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 Solution =================
// dp[u][0] = dp1: Start at u, End Deep (Type A or B)
// dp[u][1] = dp3: Start at Child, End Deep (Skip u, Type A only)
// Note: dp2 (Start u, End u) is implicitly P[u] + Sum(P[v])
struct DPState {
    long long val;
    int choice_1, choice_2; 
    // For dp1: choice_1 is the child upgraded to Deep
    // For dp3: choice_1, choice_2 are children upgraded to Deep
    int type_1; // 0: Type A (dp1), 1: Type B (dp3)
};

DPState dp[200005][2];
long long global_max_k2 = -1;
int global_root = -1;
int global_c1 = -1, global_c2 = -1;
int global_type1 = -1, global_type2 = -1; // types for global best

struct Delta {
    long long val;
    int id;
    int type; // 0: A, 1: B
};

void update_top2(vector<Delta>& top2, Delta d) {
    top2.push_back(d);
    sort(top2.begin(), top2.end(), [](const Delta& a, const Delta& b){
        return a.val > b.val;
    });
    if(top2.size() > 2) top2.pop_back();
}

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

    // Deltas for picking u
    // Allow Type A (dp1) and Type B (dp3) upgrades
    vector<Delta> diffs_u; 
    // Deltas for skipping u (dp3 for parent)
    // Allow only Type A (dp1) upgrades
    vector<Delta> diffs_skip;

    for (int v : adj[u]) {
        if (v != p) {
            long long val_a = dp[v][0].val; // dp1[v]
            long long val_b = dp[v][1].val; // dp3[v]
            
            // For picking u, we can choose best of A or B
            if (val_a >= val_b) update_top2(diffs_u, {val_a - P[v], v, 0});
            else update_top2(diffs_u, {val_b - P[v], v, 1});

            // For skipping u, only Type A allowed
            update_top2(diffs_skip, {val_a - P[v], v, 0});
        }
    }

    long long base_u = P[u] + sum_p;
    long long base_skip = sum_p;

    // dp1[u]: Pick u, upgrade 1 child (Start u, End Deep)
    dp[u][0] = {base_u, -1, -1, -1};
    if (!diffs_u.empty()) {
        if (base_u + diffs_u[0].val > base_u) {
            dp[u][0] = {base_u + diffs_u[0].val, diffs_u[0].id, -1, diffs_u[0].type};
        }
    }

    // dp3[u]: Skip u, upgrade 2 children (Start Child, End Deep, Type A only)
    dp[u][1] = {base_skip, -1, -1, -1};
    if (diffs_skip.size() >= 1 && diffs_skip[0].val > 0) {
        long long v1 = diffs_skip[0].val;
        dp[u][1].val += v1;
        dp[u][1].choice_1 = diffs_skip[0].id;
        
        if (diffs_skip.size() >= 2 && diffs_skip[1].val > 0) {
            dp[u][1].val += diffs_skip[1].val;
            dp[u][1].choice_2 = diffs_skip[1].id;
        }
    }

    // Global Max update (rooted at u)
    // Pick u, upgrade up to 2 children (Start Deep, End Deep, A or B)
    long long cur_global = base_u;
    int c1 = -1, c2 = -1, t1 = -1, t2 = -1;
    
    if (diffs_u.size() >= 1 && diffs_u[0].val > 0) {
        cur_global += diffs_u[0].val;
        c1 = diffs_u[0].id; t1 = diffs_u[0].type;
        if (diffs_u.size() >= 2 && diffs_u[1].val > 0) {
            cur_global += diffs_u[1].val;
            c2 = diffs_u[1].id; t2 = diffs_u[1].type;
        }
    }
    
    // Also consider dp3[u] (Skip u) as a candidate for global max
    // But dp3[u] is max path in subtree skipping u.
    // Usually covered by child's global max?
    // Yes, if we skip u, the path is entirely in one child's component
    // OR it connects two children components.
    // The dp3[u] calculation covers connecting two children.
    // The recursive call covers entirely in one child.
    // So just compare cur_global and dp[u][1].
    
    if (cur_global > global_max_k2) {
        global_max_k2 = cur_global;
        global_root = u;
        global_c1 = c1; global_c2 = c2;
        global_type1 = t1; global_type2 = t2;
    }
    if (dp[u][1].val > global_max_k2) {
        // This case means skipping root is better. 
        // We set a flag or just handle reconstruction carefully.
        // Actually, if skipping u is best, it implies the path is defined by
        // the two children in dp[u][1].
        // We can denote this by a special "Skip Root" state.
        global_max_k2 = dp[u][1].val;
        global_root = -u; // Negative indicates skip
        global_c1 = dp[u][1].choice_1;
        global_c2 = dp[u][1].choice_2;
    }
}

// Forward decl
void build_k2_dp1(int u, int p, vector<int>& path);
void build_k2_dp3(int u, int p, vector<int>& path);

void add_path(int v, int u, int type, vector<int>& path, bool reverse) {
    vector<int> sub;
    if (type == 0) build_k2_dp1(v, u, sub); // Type A (dp1)
    else build_k2_dp3(v, u, sub); // Type B (dp3)
    
    if (reverse) path.insert(path.end(), sub.rbegin(), sub.rend());
    else path.insert(path.end(), sub.begin(), sub.end());
}

void build_k2_dp1(int u, int p, vector<int>& path) {
    // Reconstruct dp1[u]: Pick u, upgrade 1 child
    int c = dp[u][0].choice_1;
    int t = dp[u][0].type_1;
    
    path.push_back(u);
    for (int v : adj[u]) {
        if (v != p && v != c) path.push_back(v); // Singles
    }
    
    if (c != -1) {
        add_path(c, u, t, path, false);
    }
}

void build_k2_dp3(int u, int p, vector<int>& path) {
    // Reconstruct dp3[u]: Skip u, upgrade 2 children (Type A)
    int c1 = dp[u][1].choice_1;
    int c2 = dp[u][1].choice_2;
    
    // Structure: Rev(Chain1) -> Singles -> Chain2
    // If c1 exists, add Rev(c1).
    if (c1 != -1) add_path(c1, u, 0, path, true);
    
    // Add singles (all children of u except c1, c2)
    // Note: u is SKIPPED.
    // But we need to connect c1 -> c2 via... wait.
    // If u is skipped, we assume distance allows jumping.
    // Dist v1 -> v2 is 2 (via u).
    // So valid.
    // But singles? v_single.
    // Dist c1 -> v_s is 2. v_s -> c2 is 2.
    // So yes, we can list singles.
    for (int v : adj[u]) {
        if (v != p && v != c1 && v != c2) path.push_back(v);
    }
    
    if (c2 != -1) add_path(c2, u, 0, path, false);
}

void solve_k2() {
    dfs_k2(1, 0);
    cout << global_max_k2 << "\n";
    
    vector<int> path;
    if (global_root < 0) {
        // Skip root case
        int u = -global_root;
        // Same logic as dp3 reconstruction
        int c1 = global_c1, c2 = global_c2;
        if (c1 != -1) add_path(c1, u, 0, path, true);
        for(int v : adj[u]) if(v != 0 && v != c1 && v != c2) path.push_back(v); // p=0 for root
        if (c2 != -1) add_path(c2, u, 0, path, false);
    } else {
        // Pick root case
        int u = global_root;
        int c1 = global_c1, c2 = global_c2;
        int t1 = global_type1, t2 = global_type2;
        
        // Structure: Rev(Chain1) -> u -> Singles -> Chain2
        // Note: u is picked.
        if (c1 != -1) add_path(c1, u, t1, path, true);
        
        path.push_back(u);
        for(int v : adj[u]) if(v != 0 && v != c1 && v != c2) path.push_back(v);
        
        if (c2 != -1) add_path(c2, u, t2, path, false);
    }
    
    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