제출 #1310508

#제출 시각아이디문제언어결과실행 시간메모리
1310508ayuxhkumxr22Travelling Trader (CCO23_day2problem2)C++20
0 / 25
7 ms13120 KiB
#include <iostream>
#include <vector>
#include <algorithm>
#include <tuple>

using namespace std;

const int MAXN = 200005;
const int MAXK = 4; // K is at most 3

int N, K;
int P[MAXN];
vector<int> adj[MAXN];

// DP Table: dp[u][s_in][s_out] = max_profit
// Since s_out is small, we return a fixed size array.
// -1 indicates impossible state.
struct Result {
    int val[MAXK]; // val[s_out] = max_profit
    Result() {
        fill(val, val + MAXK, -1);
    }
};

Result dp[MAXN][MAXK];
bool visited_dp[MAXN][MAXK];

// Reconstruction structures
// Store decisions: type (0: stop, 1: pick U, 2: enter child)
struct Decision {
    int type; 
    int next_slack; // slack passed to next step
    int child_exit; // if entering child, what was child's exit slack
    int gained_here; // profit gained in this step
    bool u_picked; // new u_picked status
};

// path[u][s_in][step][current_slack][u_picked_status]
// step: 0 to adj.size()
// Map isn't efficient, but logic is complex. 
// We will rebuild solution by re-running logic with "finding" target.
// To save memory, we won't store full table, but just re-run logic for reconstruction.

// Helper to update max
void update(int& target, int candidate) {
    if (candidate > target) target = candidate;
}

// Compute DP for node u
void solve(int u, int p) {
    // Precompute children
    vector<int> children;
    for (int v : adj[u]) {
        if (v != p) {
            solve(v, u); // Process children first
            children.push_back(v);
        }
    }

    // For each possible entry slack from parent
    for (int s_in = 0; s_in <= K; s_in++) {
        Result& res = dp[u][s_in];
        
        // Local DP State: current[slack][u_picked] = profit
        int cur[MAXK][2];
        for(int s=0; s<=K; ++s) { cur[s][0] = -1; cur[s][1] = -1; }

        // Initial State at Slot 0 (Pre-order)
        // We arrived at u with s_in slack. 
        // Note: s_in is slack measured AT u.
        // Option A: Just arrived, haven't picked u.
        cur[s_in][0] = 0;
        
        // Option B: Pick u immediately (if slack allows, i.e., we reached here)
        // If s_in >= 0, we are alive. Picking u resets slack to K.
        if (s_in >= 0) {
             cur[K][1] = P[u];
        }

        // Capture results from "stopping early" (returning to parent immediately)
        // If we return to parent immediately from state (sl, picked), 
        // distance u->p is 1. Parent sees slack (sl - 1).
        for(int s=0; s<=K; ++s) {
            if (cur[s][0] != -1 && s-1 >= 0) update(res.val[s-1], cur[s][0]);
            if (cur[s][1] != -1 && s-1 >= 0) update(res.val[s-1], cur[s][1]);
        }

        // Iterate through children
        for (int v : children) {
            int next_cur[MAXK][2];
            for(int s=0; s<=K; ++s) { next_cur[s][0] = -1; next_cur[s][1] = -1; }

            // Try to transition from current states
            for (int sl = 0; sl <= K; sl++) {
                for (int picked = 0; picked <= 1; picked++) {
                    int val = cur[sl][picked];
                    if (val == -1) continue;

                    // Option 1: Enter child v
                    // Cost to enter: 1 step. Entry slack at v: sl - 1.
                    if (sl - 1 >= 0) {
                        int child_in = sl - 1;
                        // Check child's outcomes
                        for (int child_out = 0; child_out <= K; child_out++) {
                            int gain = dp[v][child_in].val[child_out];
                            if (gain == -1) continue;

                            // Return from child: 1 step. 
                            // Slack at u becomes child_out - 1.
                            if (child_out - 1 >= 0) {
                                int new_sl = child_out - 1;
                                update(next_cur[new_sl][picked], val + gain);
                            }
                        }
                    }
                    
                    // Option 2: Pick u "between" children (if not picked)
                    // This is handled by processing the 'cur' state before moving to next child
                    // Wait, picking U is a Slot action.
                    // My loop is Child -> Child.
                    // The "Slot" actions can be merged. 
                    // We can pick U "after" returning from child v, before v_next.
                    // This is equivalent to transitioning `next_cur` states.
                }
            }
            
            // Update 'cur' to 'next_cur' (results after traversing child)
            for(int s=0; s<=K; ++s) {
                cur[s][0] = next_cur[s][0];
                cur[s][1] = next_cur[s][1];
            }
            
            // After traversing child, we are at a "Slot".
            // 1. We can choose to pick U now (if not picked).
            // 2. We can choose to Stop and return to parent.
            for (int sl = 0; sl <= K; sl++) {
                // If we haven't picked u, we can pick it now
                if (cur[sl][0] != -1) {
                    update(cur[K][1], cur[sl][0] + P[u]);
                }
                
                // Record results if we stop here
                if (cur[sl][0] != -1 && sl - 1 >= 0) update(res.val[sl-1], cur[sl][0]);
                if (cur[sl][1] != -1 && sl - 1 >= 0) update(res.val[sl-1], cur[sl][1]);
            }
        }
    }
}

// Global path storage
vector<int> path;

// Re-run the logic to find the path
void reconstruct(int u, int p, int s_in, int target_s_out, int target_profit) {
    vector<int> children;
    for (int v : adj[u]) if (v != p) children.push_back(v);

    // We need to find the sequence of decisions.
    // We run the same DP logic but break when we find the matching intermediate state.
    
    // States are now specific: we know we MUST reach (target_s_out, target_profit).
    // We backtrack from the end.
    
    // Stores history: history[child_idx][slack][picked] = {prev_slack, prev_picked, profit_at_step}
    // This is hard because of the "max" overwriting.
    // Easier approach: Recursive search.
    
    // Let's try to match the result.
    
    // 1. Did we stop at the very end (after last child)?
    //    Or did we stop earlier?
    //    Or did we pick U at the very end?
    
    // Re-simulate forward, keep parent pointers?
    // Given N=200,000, K=3, full DP table in memory is fine.
    // Let's store the `cur` tables for each step.
    
    // steps: 0 (pre), 1 (after child 1), ..., k (after child k)
    // We need `vector<vector<pair<int,int>>> traces`
    // Memory: N nodes * degree * K states. Sum of degrees is 2N. Total memory O(NK). Safe.
    
    struct StepInfo {
        int profit;
        int prev_sl;
        int prev_picked;
        int action; // 0: pass, 1: pick U, 2: traverse child
        int child_exit; // if traversed child
    };
    
    // DP Table for this specific (u, s_in): 
    // table[step][slack][picked] = StepInfo
    // We re-compute this table on demand.
    
    int num_steps = children.size();
    // dim: [step+1][K+1][2]
    // using vector to be dynamic
    vector<vector<vector<StepInfo>>> table(num_steps + 1, vector<vector<StepInfo>>(K + 1, vector<StepInfo>(2, {-2, -1, -1, -1, -1})));
    
    // Init step 0
    table[0][s_in][0] = {0, -1, -1, 0, -1}; // Start state
    if (s_in >= 0) table[0][K][1] = {P[u], s_in, 0, 1, -1}; // Pick U immediately
    
    for (int i = 0; i < num_steps; ++i) {
        int v = children[i];
        
        // Transition from table[i] to table[i+1]
        for (int sl = 0; sl <= K; sl++) {
            for (int picked = 0; picked <= 1; picked++) {
                if (table[i][sl][picked].profit == -2) continue;
                int curr_prof = table[i][sl][picked].profit;
                
                // Action: Enter child
                if (sl - 1 >= 0) {
                    for (int c_out = 0; c_out <= K; c_out++) {
                        int gain = dp[v][sl-1].val[c_out];
                        if (gain != -1 && c_out - 1 >= 0) {
                            int n_sl = c_out - 1;
                            int n_prof = curr_prof + gain;
                            if (n_prof > table[i+1][n_sl][picked].profit) {
                                table[i+1][n_sl][picked] = {n_prof, sl, picked, 2, c_out};
                            }
                        }
                    }
                }
            }
        }
        
        // After transition, allow picking U
        for (int sl = 0; sl <= K; sl++) {
            if (table[i+1][sl][0].profit != -2) {
                int base = table[i+1][sl][0].profit;
                int n_prof = base + P[u];
                if (n_prof > table[i+1][K][1].profit) {
                    // We mark this as Action 1 (Pick U) occurring at step i+1
                    // Predecessor is the same step i+1, state (sl, 0)
                    // We encode this carefully. 
                    // Actually, let's process "Pick U" as a separate micro-step or update in place.
                    // Simple way: Update in place, but store that we came from (sl, 0)
                    // We'll mark prev_sl as sl, prev_picked as 0, action as 1.
                    // But we are overwriting table[i+1][K][1].
                    // Correct logic: The state K,1 is reached from sl,0 at same step.
                    table[i+1][K][1] = {n_prof, sl, 0, 1, -1};
                }
            }
        }
    }
    
    // Now Trace Back from 'target'
    // We need to find WHICH step/state produced the (target_s_out, target_profit) result.
    // The result was formed by returning to parent (slack - 1).
    // So we look for a state in 'table' where (slack - 1) == target_s_out.
    // And profit == target_profit.
    // It could be at any step (0 to num_steps).
    
    int end_step = -1, end_sl = -1, end_picked = -1;
    
    for (int i = 0; i <= num_steps; ++i) {
        for (int sl = 0; sl <= K; sl++) {
            for (int p = 0; p <= 1; p++) {
                if (table[i][sl][p].profit == target_profit && sl - 1 == target_s_out) {
                    end_step = i; end_sl = sl; end_picked = p;
                    goto found;
                }
            }
        }
    }
    found:;
    
    // Backtrack path
    vector<Decision> trace;
    int cur_step = end_step, cur_sl = end_sl, cur_picked = end_picked;
    
    while (true) {
        if (cur_step == 0 && cur_sl == s_in && cur_picked == 0) break; // Start
        
        StepInfo info = table[cur_step][cur_sl][cur_picked];
        
        if (info.action == 1) { // Picked U here
            trace.push_back({1, 0, 0, 0, 0}); // Marker to pick U
            // Go to prev state (same step)
            cur_sl = info.prev_sl;
            cur_picked = info.prev_picked;
        } else if (info.action == 2) { // Traversed Child
            // This transition was from step-1 to step
            trace.push_back({2, 0, info.child_exit, 0, 0}); // Marker child
            cur_step--;
            cur_sl = info.prev_sl;
            cur_picked = info.prev_picked;
        } else if (info.action == 0) { // Just Start step 0 pick U
             // Handled by loop condition? 
             // If action 0, it must be the "Pick U at start" case
             trace.push_back({1, 0, 0, 0, 0});
             // Base was s_in, 0
             cur_sl = info.prev_sl;
             cur_picked = info.prev_picked;
        }
    }
    
    reverse(trace.begin(), trace.end());
    
    // Execute Trace
    // Reconstruct order: 
    // Start.
    // trace[0] might be Pick U.
    // Then children...
    
    int t_idx = 0;
    // Check if initial Pick U happened
    if (t_idx < trace.size() && trace[t_idx].type == 1) {
        path.push_back(u);
        t_idx++;
    }
    
    for (int i = 0; i < num_steps; ++i) {
        // Child i
        if (t_idx < trace.size() && trace[t_idx].type == 2) {
            // We entered child i
            // We need to calc correct parameters to recurse
            // We need the entry slack into child.
            // We need the child's profit contribution.
            // This requires storing more in 'StepInfo' or recalculating.
            // Recalculating:
            // The step transition was: 
            // table[i][prev_sl][prev_p] + child -> table[i+1]...
            // We know the child exit slack (stored in Decision).
            // We know child entry slack = prev_sl - 1.
            // We need child profit. 
            // Child profit = total_now - total_prev.
            // But we don't have total_prev easily accessible in loop?
            // Actually, we do. We iterate forward.
            
            // Wait, tracking 'current state' in forward pass is hard without profit values.
            // But we have the 'trace' which tells us what happened.
            // We can retrieve the profit from the DP table again.
            
            // Correct approach:
            // We are at step i. Trace says "Traverse child".
            // We need to look at table[i] state to get slack.
            // But 'table' is local to this function.
            // We have it!
            
            // Wait, 'trace' doesn't store the state indices, just the action type.
            // We need to maintain the state as we walk forward.
            // Let's re-simulate.
        }
    }
    
    // Better Forward Execution:
    int c_sl = s_in;
    int c_pk = 0;
    int child_idx = 0;
    
    for (const auto& dec : trace) {
        if (dec.type == 1) {
            path.push_back(u);
            c_sl = K; 
            c_pk = 1;
        } else if (dec.type == 2) {
            // Enter child
            int v = children[child_idx];
            int entry_s = c_sl - 1;
            int exit_s = dec.child_exit;
            int child_profit = dp[v][entry_s].val[exit_s];
            
            reconstruct(v, u, entry_s, exit_s, child_profit);
            
            c_sl = exit_s - 1;
            child_idx++;
        }
    }
}


int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    if (!(cin >> N >> K)) return 0;
    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];
    }

    // Solve DP
    // Start at node 1.
    // Entry slack?
    // Problem: "trader starts by doing business in city 1".
    // This implies we MUST pick 1 first.
    // And distance to next node is relative to 1.
    // So essentially, we enter 1 with enough slack to pick it, 
    // and we are FORCED to pick it.
    
    // We can simulate this by calling solve(1, 0).
    // And looking at the result where 1 is picked.
    // But solve() computes best strategy.
    // We need to extract the specific result where 1 is picked.
    // However, our DP stores results for all entry slacks.
    // If we say entry slack is K (fresh), and we check results.
    // Wait, the "start" is outside the recursion.
    // Root has entry slack... say K? Or 0?
    // If we assume the "virtual previous node" was at distance 1?
    // Actually, we can just look at dp[1][K].
    // And ensure 1 is picked?
    // Our DP returns max profit. Does it guarantee 1 is picked?
    // Not necessarily. But since p_i >= 1, picking 1 is always better than skipping it 
    // if we are at the start and have full slack (unless it prevents a huge subtree).
    // But since K is small and we start AT 1, we simply pick 1.
    // Slack becomes K.
    // So we basically want to enter 1, pick it, and see what happens.
    // This corresponds to Option B in "Slot 0" logic inside solve().
    // We can just use the DP table result for s_in = 0 (or K).
    // And finding the max profit among valid s_out.
    // Actually, since we start at 1, we effectively arrive at 1 with slack K?
    // If we arrive with K, we pick 1 -> slack K.
    // If we don't pick 1 -> slack K? No.
    // The problem forces picking 1.
    // So we want the reconstruction to force picking 1.
    
    solve(1, 0);

    // Find best total profit starting at 1
    // We can assume we enter 1 with "virtual slack" that allows picking.
    // Let's use s_in = K.
    // Since we MUST pick 1, we look for the max profit solution in dp[1][K]
    // where the path effectively starts with 1.
    // Actually, our DP maximizes profit. If picking 1 is optimal, it does it.
    // If skipping 1 was optimal (for the general recursion), it might skip.
    // BUT, problem constraint: "starts by doing business in city 1".
    // So we must force it.
    // My DP doesn't have a "must pick u" flag passed in.
    // However, since we reconstructed, we can check.
    // Or, we can just manually pick 1 and then process children.
    // "Process children with slack K-1 and sum their results"?
    // No, we can bounce between children.
    
    // Simplest Fix:
    // Run reconstruction on dp[1][K], but FORCE the first decision to be "Pick 1".
    // Or check if dp[1][K] picked 1.
    // Wait, if picking 1 is mandatory, and my DP didn't pick it, my DP value is wrong for the problem.
    // Is it possible that skipping 1 yields more profit? 
    // Maybe, to reach a deep node?
    // But we are AT 1. Picking it costs 0 distance. It resets slack to K.
    // It is ALWAYS better or equal to pick 1 than to arrive at 1 and skip it (leaving with same slack).
    // Picking 1 sets slack to K (max possible). Skipping leaves slack at K (arrival) or less.
    // Since K >= s_arrival, picking 1 increases slack AND profit.
    // So the optimal strategy WILL pick 1.
    
    int best_profit = -1;
    int best_s_out = -1;
    for(int s=0; s<=K; ++s) {
        if(dp[1][K].val[s] > best_profit) {
            best_profit = dp[1][K].val[s];
            best_s_out = s;
        }
    }
    
    cout << best_profit << "\n";
    
    reconstruct(1, 0, K, best_s_out, best_profit);
    
    cout << path.size() << "\n";
    for (int i = 0; i < path.size(); i++) {
        cout << path[i] << (i == path.size() - 1 ? "" : " ");
    }
    cout << "\n";

    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...