제출 #1310526

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1310526		ayuxhkumxr22	Travelling Trader (CCO23_day2problem2)	C++20	25 / 25	122 ms	50364 KiB

Main

/*
 * Problem: Travelling Trader (CCO 2023 P2)
 * Logic: Ported from correct C solution.
 * K=1: DFS Longest Path.
 * K=3: DFS Euler Tour.
 * K=2: 3-State Tree DP with "Greedy Neighbors" and "Abandonment".
 */

#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 Logic
// ==========================================
long long dp_k1[200005];
int best_child_k1[200005];

void dfs1_calc(int u, int p) {
    long long max_val = 0;
    int bc = -1;
    for (int v : adj[u]) {
        if (v != p) {
            dfs1_calc(v, u);
            if (dp_k1[v] > max_val) {
                max_val = dp_k1[v];
                bc = v;
            }
        }
    }
    dp_k1[u] = max_val + P[u];
    best_child_k1[u] = bc;
}

void dfs1_build(int u, vector<int>& path) {
    path.push_back(u);
    if (best_child_k1[u] != -1) dfs1_build(best_child_k1[u], path);
}

void solve_k1() {
    fill(best_child_k1, best_child_k1 + N + 1, -1);
    dfs1_calc(1, 0);
    
    vector<int> path;
    dfs1_build(1, path);
    
    cout << dp_k1[1] << "\n";
    cout << path.size() << "\n";
    for(int i=0; i<path.size(); ++i) cout << path[i] << (i==path.size()-1?"":" ");
    cout << "\n";
}

// ==========================================
// K = 3: Euler Tour Logic
// ==========================================
void dfs3_build(int u, int p, bool rev, vector<int>& path) {
    if (!rev) path.push_back(u);
    for (int v : adj[u]) {
        if (v != p) dfs3_build(v, u, rev ^ 1, path);
    }
    if (rev) path.push_back(u);
}

void solve_k3() {
    long long total = 0;
    for(int i=1; i<=N; ++i) total += P[i];
    
    vector<int> path;
    dfs3_build(1, 0, false, path);
    
    cout << total << "\n";
    cout << path.size() << "\n";
    for(int i=0; i<path.size(); ++i) cout << path[i] << (i==path.size()-1?"":" ");
    cout << "\n";
}

// ==========================================
// K = 2: Complex Tree DP (The Hard Part)
// ==========================================
// Direct mapping of C code variables:
// dp[i] -> dp_tail[i] (Best path ending deep)
// dq[i] -> dp_ans[i]  (Best valid component in subtree)
// dr[i] -> dp_mix[i]  (Complex intermediate state)
// ww[i] -> P[i]

long long dp_tail[200005], dp_ans[200005], dp_mix[200005];
// Store "best children" for reconstruction
// jjp1, jjp2, jjp3 -> best_tail_1, best_tail_2, best_tail_3
int best_tail_1[200005], best_tail_2[200005], best_tail_3[200005];
// jjq, ttq -> best_ans_child, type_ans (2 or 3)
int best_ans_child[200005], type_ans[200005];
// jjr, ttr -> best_mix_child, type_mix (2 or 3)
int best_mix_child[200005], type_mix[200005];

void dfs2_calc(int u, int p) {
    // 1. Calculate 'w': Sum of all children profits
    long long w = 0;
    for (int v : adj[u]) if (v != p) {
        dfs2_calc(v, u);
        w += P[v];
    }

    // 2. Calculate dp_tail[u] (and find top 3 children by dp_tail)
    // dp[i] = max(dp[child]) + w
    int jp1 = -1, jp2 = -1, jp3 = -1;
    for (int v : adj[u]) if (v != p) {
        if (jp1 == -1 || dp_tail[jp1] < dp_tail[v]) {
            jp3 = jp2; jp2 = jp1; jp1 = v;
        } else if (jp2 == -1 || dp_tail[jp2] < dp_tail[v]) {
            jp3 = jp2; jp2 = v;
        } else if (jp3 == -1 || dp_tail[jp3] < dp_tail[v]) {
            jp3 = v;
        }
    }
    dp_tail[u] = (jp1 == -1 ? 0 : dp_tail[jp1]) + w;
    best_tail_1[u] = jp1; best_tail_2[u] = jp2; best_tail_3[u] = jp3;

    // 3. Calculate dp_ans[u] (dq)
    long long best_val = -1; 
    int best_c = -1, type = 2;

    for (int v : adj[u]) if (v != p) {
        // Option 2 (Merge): dp[other_tail] + dq[v] + w
        int other = (v != jp1 ? jp1 : jp2);
        long long val_opt2 = (other == -1 ? 0 : dp_tail[other]) + dp_ans[v] + w;
        if (best_val < val_opt2) {
            best_val = val_opt2; best_c = v; type = 2;
        }

        // Option 3 (Abandon): dr[v] + P[v] (Note: w is NOT added!)
        long long val_opt3 = dp_mix[v] + P[v];
        if (best_val < val_opt3) {
            best_val = val_opt3; best_c = v; type = 3;
        }
    }
    dp_ans[u] = (best_val == -1 ? 0 : best_val);
    best_ans_child[u] = best_c; type_ans[u] = type;

    // 4. Calculate dp_mix[u] (dr)
    best_val = -1; best_c = -1; type = 2;
    for (int v : adj[u]) if (v != p) {
        // Identify best 2 tails excluding v
        int o1, o2;
        if (v == jp1) { o1 = jp2; o2 = jp3; }
        else if (v == jp2) { o1 = jp1; o2 = jp3; }
        else { o1 = jp1; o2 = jp2; }

        // Option 2 (Merge Mix): dp[o1] + dp[o2] + dq[v] + w
        long long val_opt2 = (o1 == -1 ? 0 : dp_tail[o1]) + (o2 == -1 ? 0 : dp_tail[o2]) + dp_ans[v] + w;
        if (best_val < val_opt2) {
            best_val = val_opt2; best_c = v; type = 2;
        }

        // Option 3 (Deep Mix): dp[o1] + dr[v] + w
        long long val_opt3 = (o1 == -1 ? 0 : dp_tail[o1]) + dp_mix[v] + w;
        if (best_val < val_opt3) {
            best_val = val_opt3; best_c = v; type = 3;
        }
    }
    dp_mix[u] = (best_val == -1 ? 0 : best_val);
    best_mix_child[u] = best_c; type_mix[u] = type;
}

// Reconstruction Logic (Mapping C dfs2_ modes)
// mode -1: Build dp_tail (Reverse: Path -> Tail)
// mode 1:  Build dp_tail (Forward: Tail -> Path) ? No, check code.
// C code logic for modes:
// -1: Tail logic. Print neighbors != jp1. Recurse jp1 (mode 1). Print u.
//  1: Reverse Tail. Print u. Recurse jp1 (mode -1). Print neighbors != jp1.
//  2: Root (dq). Print u. Recurse jp1 (mode -1). Print neighbors. Recurse jq (mode 2/3).
//  3: Mix (dr). Print... depends on ttr.

vector<int> path_k2;

void dfs2_build(int u, int p, int mode) {
    // Helper to add all neighbors except specific ones
    auto add_neighbors = [&](int ex1, int ex2, int ex3) {
        for (int v : adj[u]) if (v != p && v != ex1 && v != ex2 && v != ex3) {
            path_k2.push_back(v);
        }
    };

    if (mode == -1) { 
        // Logic: Neighbors -> Recurse Tail (1) -> u
        int jp1 = best_tail_1[u];
        add_neighbors(jp1, -1, -1);
        if (jp1 != -1) dfs2_build(jp1, u, 1);
        path_k2.push_back(u);
    } 
    else if (mode == 1) {
        // Logic: u -> Recurse Tail (-1) -> Neighbors
        int jp1 = best_tail_1[u];
        path_k2.push_back(u);
        if (jp1 != -1) dfs2_build(jp1, u, -1);
        add_neighbors(jp1, -1, -1);
    } 
    else if (mode == 2) { // Building dp_ans (dq)
        int jq = best_ans_child[u];
        // Which tail was used?
        int jp1 = (jq != best_tail_1[u] ? best_tail_1[u] : best_tail_2[u]);
        
        if (type_ans[u] == 2) {
            // Option 2: Tail + Neighbors + dq[child]
            // u -> Recurse Tail (-1) -> Neighbors -> Recurse dq (2)
            path_k2.push_back(u);
            if (jp1 != -1) dfs2_build(jp1, u, -1);
            add_neighbors(jp1, jq, -1);
            if (jq != -1) dfs2_build(jq, u, 2);
        } else {
            // Option 3 (Abandon): u -> Recurse dr (3)
            // Note: Neighbors are SKIPPED here!
            path_k2.push_back(u);
            dfs2_build(jq, u, 3);
        }
    } 
    else { // mode 3 (Building dp_mix / dr)
        int jr = best_mix_child[u];
        // Identify tails used
        int jp1, jp2;
        if (jr == best_tail_1[u]) { jp1 = best_tail_2[u]; jp2 = best_tail_3[u]; }
        else if (jr == best_tail_2[u]) { jp1 = best_tail_1[u]; jp2 = best_tail_3[u]; }
        else { jp1 = best_tail_1[u]; jp2 = best_tail_2[u]; }

        if (type_mix[u] == 2) {
            // Option 2: Tail1 (1) -> u -> Tail2 (-1) -> Neighbors -> dq (2)
            // Wait, C code order: 
            // Recurse jp1 (1). u. Recurse jp2 (-1). Neighbors. Recurse jr (2).
            if (jp1 != -1) dfs2_build(jp1, u, 1);
            path_k2.push_back(u);
            if (jp2 != -1) dfs2_build(jp2, u, -1);
            add_neighbors(jp1, jp2, jr);
            if (jr != -1) dfs2_build(jr, u, 2);
        } else {
            // Option 3: Neighbors -> Recurse Tail1 (1) -> u -> Recurse dr (3)
            add_neighbors(jp1, jr, -1);
            if (jp1 != -1) dfs2_build(jp1, u, 1);
            path_k2.push_back(u);
            dfs2_build(jr, u, 3);
        }
    }
}

void solve_k2() {
    dfs2_calc(1, 0);
    // Final answer is dq[1] + P[1]? No, C code says dq[0] + ww[0].
    // Note: C code uses 0-based indexing. P[u] needs to be added manually at root?
    // My dp_ans logic seems to sum children vals. P[root] is separate.
    // Yes, dp_ans[u] sums children contributions. P[u] is separate.
    cout << dp_ans[1] + P[1] << "\n";
    
    // Build path starting with mode 2 (Root Answer)
    dfs2_build(1, 0, 2);
    
    cout << path_k2.size() << "\n";
    for(int i=0; i<path_k2.size(); ++i) cout << path_k2[i] << (i==path_k2.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: 2 / 2
• Subtask 2: 7 / 7
• Subtask 3: 3 / 3
• Subtask 4: 4 / 4
• Subtask 5: 4 / 4
• Subtask 6: 5 / 5