Submission #1309567

#	Time	Username	Problem	Language	Result	Execution time	Memory
1309567		ayuxhkumxr22	Travelling Trader (CCO23_day2problem2)	C++20	0 / 25	7 ms	1092 KiB

Main

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>

using namespace std;

// --- Constants & Globals ---
const int MAXN = 200005;
const long long INF = 1e18;

int N, K;
vector<int> adj[MAXN];
vector<int> tree[MAXN]; // Directed tree structure after DFS
long long P[MAXN];      // Profit of city i

// DP State: dp[u][d_in] stores a list of possible outcomes {d_out, profit}
// d_in: distance from the nearest selected ancestor (0..K)
// d_out: distance to the nearest selected descendant (0..K)
struct State {
    int d_out;
    long long profit;
};

// Memoization table
// d_in can range from 0 to K. We use size 4 since K <= 3.
vector<State> dp[MAXN][4];
bool visited[MAXN][4];

// --- Helper Functions ---

// Filter Pareto optimal states to keep the list small
// We want small d_out and large profit.
// If state A has (d_out <= B.d_out) and (profit >= B.profit), A dominates B.
void optimize_states(vector<State>& states) {
    if (states.empty()) return;
    
    // Sort by d_out ascending, then profit descending
    sort(states.begin(), states.end(), [](const State& a, const State& b) {
        if (a.d_out != b.d_out) return a.d_out < b.d_out;
        return a.profit > b.profit;
    });

    vector<State> optimal;
    long long max_p = -1;
    
    for (const auto& s : states) {
        if (s.profit > max_p) {
            optimal.push_back(s);
            max_p = s.profit;
        }
    }
    states = optimal;
}

// Pre-processing to build tree direction
void build_tree(int u, int p) {
    for (int v : adj[u]) {
        if (v != p) {
            tree[u].push_back(v);
            build_tree(v, u);
        }
    }
}

// --- Core DP Logic ---

// Returns the best state (max profit) found by the Greedy Merge strategy
// picking_u: boolean, true if we select current node u
// d_in_start: the initial gap. If picking_u is true, this is effectively 0.
State get_greedy_outcome(int u, int d_in_start, bool picking_u) {
    long long current_profit = (picking_u ? P[u] : 0);
    int current_gap = (picking_u ? 0 : d_in_start);
    
    // If the starting gap is already invalid, return invalid state
    // Note: If picking_u is true, current_gap is 0, always valid.
    // If skipping, we must check.
    if (current_gap >= K && !picking_u) return {K + 1, -1}; 

    // We keep track of which children are used to avoid cycles in our greedy selection
    vector<bool> child_used(tree[u].size(), false);
    
    // Greedy Simulation:
    // Repeatedly try to visit a child that accepts (current_gap + 1).
    // Among candidates, pick the one that returns the smallest d_out (to keep the path alive).
    // Tie-break with max profit.
    while (true) {
        int req_in = current_gap + 1;
        if (req_in > K) break; // Cannot extend further

        int best_child_idx = -1;
        int best_state_idx = -1;
        int best_next_gap = 100; // Minimize this
        long long best_gain = -1; // Maximize this

        // Scan all children to find the best candidate
        for (int i = 0; i < tree[u].size(); i++) {
            if (child_used[i]) continue;
            int v = tree[u][i];

            // Look at the child's DP results for input req_in
            // We assume dp[v] is already computed
            for (int s_i = 0; s_i < dp[v][req_in].size(); s_i++) {
                State& s = dp[v][req_in][s_i];
                // If we visit this child, the new gap at u becomes (child's d_out + 1)
                int next_g = s.d_out + 1;

                // Priority: Lower Gap > Higher Profit
                if (next_g < best_next_gap) {
                    best_next_gap = next_g;
                    best_gain = s.profit;
                    best_child_idx = i;
                    best_state_idx = s_i;
                } else if (next_g == best_next_gap) {
                    if (s.profit > best_gain) {
                        best_gain = s.profit;
                        best_child_idx = i;
                        best_state_idx = s_i;
                    }
                }
            }
        }

        // If we found a valid move, take it
        if (best_child_idx != -1) {
            child_used[best_child_idx] = true;
            current_profit += best_gain;
            current_gap = best_next_gap;
        } else {
            break; // No more children can be added
        }
    }
    
    return {current_gap, current_profit};
}

void solve(int u) {
    // Solve for all children first (Post-order)
    for (int v : tree[u]) {
        solve(v);
    }

    // Determine DP states for u
    // We calculate valid d_in from 0 to K (actually K is strictly limit, so inputs 0..K-1 are relevant?
    // d_in can be K if we don't pick u, but then we can't extend.
    // The problem says "more than K days", so gap K is allowed (distance K+1 is bad).
    // So d_in <= K is valid input.
    
    // --- Option 1: Pick u ---
    // This resets the gap to 0. The outcome is the same regardless of d_in (as long as d_in allows reaching u).
    State picked_res = get_greedy_outcome(u, 0, true);

    // --- Option 2: Skip u ---
    // We calculate this for each possible input d_in
    for (int d = 0; d <= K; d++) {
        // If we pick u, it covers all d inputs (since it resets)
        dp[u][d].push_back(picked_res);

        // If we skip u
        if (d < K) { // Can only skip if we don't violate K immediately
             // d_in passed to greedy is d. It will try to pass d+1 to children.
             State skipped_res = get_greedy_outcome(u, d, false);
             if (skipped_res.profit != -1) {
                 dp[u][d].push_back(skipped_res);
             }
        }
        
        // Prune the list
        optimize_states(dp[u][d]);
    }
}

// --- Reconstruction ---
vector<int> final_path;

void reconstruct(int u, int d_in, bool parent_picked) {
    // We need to determine if we picked u or skipped u.
    // We re-run the logic to see which option yields the optimal profit recorded in DP.
    
    // Find expected profit from DP table
    long long target_profit = -1;
    // We need the BEST profit from the pareto list that fits our context?
    // Actually, the caller ensures we are on the optimal path.
    // But dp[u][d_in] has multiple options (tradeoff between profit and d_out).
    // However, since we start at root with d_in=0 and pick u, we just maximize profit.
    // For recursion, we need to know WHICH state was chosen by the parent greedy.
    // This is complex.
    
    // Alternative: The greedy simulation in the parent determined:
    // 1. Which child to visit
    // 2. Which state index of that child used
    // We should pass "Target Profit" or specific decision down?
    // Let's re-simulate the parent's greedy choice to find the child's requirements.
}

// Recursive reconstruction that re-runs the greedy logic
void reconstruct_greedy(int u, int d_in, bool picking_u) {
    if (picking_u) final_path.push_back(u);
    
    int current_gap = (picking_u ? 0 : d_in);
    
    // Re-run the exact greedy loop to find path
    vector<bool> child_used(tree[u].size(), false);
    
    while (true) {
        int req_in = current_gap + 1;
        if (req_in > K) break; 

        int best_child_idx = -1;
        int best_state_idx = -1;
        int best_next_gap = 100;
        long long best_gain = -1;

        for (int i = 0; i < tree[u].size(); i++) {
            if (child_used[i]) continue;
            int v = tree[u][i];

            for (int s_i = 0; s_i < dp[v][req_in].size(); s_i++) {
                State& s = dp[v][req_in][s_i];
                int next_g = s.d_out + 1;
                if (next_g < best_next_gap) {
                    best_next_gap = next_g;
                    best_gain = s.profit;
                    best_child_idx = i;
                    best_state_idx = s_i;
                } else if (next_g == best_next_gap) {
                    if (s.profit > best_gain) {
                        best_gain = s.profit;
                        best_child_idx = i;
                        best_state_idx = s_i;
                    }
                }
            }
        }

        if (best_child_idx != -1) {
            child_used[best_child_idx] = true;
            // Recurse: We must determine if the child was picked or skipped.
            // We know the child achieved 'best_gain' with output 'best_next_gap-1'.
            // We check if "Pick" or "Skip" at child generates this.
            int v = tree[u][best_child_idx];
            
            // Check Pick
            State pick_res = get_greedy_outcome(v, 0, true);
            bool is_pick = false;
            
            // Note: Pareto optimization might remove exact matches, but we kept optimal ones.
            // If Pick matches the stats:
            // However, Skip might also match. Tie-breaking rules in solve must match here.
            // In optimize_states, we prefer picking if profits equal? No specific order.
            // Let's check if Pick provides the profit/d_out.
            
            // Actually, simply:
            // Calculate Pick Result at child.
            // Calculate Skip Result at child (with input req_in).
            // Compare which one matches 'best_gain' and 'best_next_gap - 1'.
            // If both match, prefer Pick (arbitrary, but must be consistent with dp build).
            // In build, we added Pick then Skip. Optimize sort stable? No.
            // But usually Pick has better profit.
            
            State skip_res = {100, -1};
            if (req_in < K) skip_res = get_greedy_outcome(v, req_in, false);
            
            // Check match
            // We use >= for profit to capture the "best" choice logic
            if (pick_res.d_out == best_next_gap - 1 && pick_res.profit == best_gain) {
                is_pick = true;
            } else if (skip_res.profit != -1 && skip_res.d_out == best_next_gap - 1 && skip_res.profit == best_gain) {
                is_pick = false;
            } else {
                // Fallback: If optimized out, usually picking is the robust choice for reconstruction
                is_pick = true; 
            }
            
            reconstruct_greedy(v, req_in, is_pick);
            
            current_gap = best_next_gap;
        } else {
            break;
        }
    }
}

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

    // Root is 1
    build_tree(1, 0);

    solve(1);

    // Initial call: The problem implies we start at City 1 and do business.
    // So we MUST pick node 1.
    State res = dp[1][0][0]; // Best result for picking 1 (input 0)
    
    // Find the specific result that gave max profit.
    // Since we forced Pick 1, we look at get_greedy_outcome(1, 0, true)
    
    cout << res.profit << endl;
    
    reconstruct_greedy(1, 0, true);
    
    cout << final_path.size() << endl;
    for (int i = 0; i < final_path.size(); i++) {
        cout << final_path[i] << (i == final_path.size() - 1 ? "" : " ");
    }
    cout << endl;

    return 0;
}

• Subtask 1: 0 / 2
• Subtask 2: 0 / 7
• Subtask 3: 0 / 3
• Subtask 4: 0 / 4
• Subtask 5: 0 / 4
• Subtask 6: 0 / 5