Submission #1310509

#TimeUsernameProblemLanguageResultExecution timeMemory
1310509ayuxhkumxr22Travelling Trader (CCO23_day2problem2)C++20
Compilation error
0 ms0 KiB
#include <iostream>
#include <vector>
#include <algorithm>
#include <tuple>

using namespace std;

// Constraints
const int MAXN = 200005;
const int MAXK = 3; 

// Input data
int N, K;
long long P[MAXN];
vector<int> adj[MAXN];

// Constants for invalid states
const long long INF_LL = -1e18; 

// DP State structure
struct State {
    long long closed[MAXK][MAXK]; // [first_dist][last_dist]
    long long open[MAXK];         // [first_dist]
    long long global_max;         // Best internal path

    State() {
        for(int i=0; i<MAXK; ++i) {
            open[i] = INF_LL;
            for(int j=0; j<MAXK; ++j) {
                closed[i][j] = INF_LL;
            }
        }
        global_max = INF_LL;
    }
};

// Backtracking structures
enum SourceType : char {
    NONE = 0,
    FROM_L = 1,      // Result came from Left state alone
    FROM_R = 2,      // Result came from Right state alone
    L_THEN_R = 3,    // Left chain -> Right chain
    R_THEN_L = 4     // Right chain -> Left chain
};

struct TraceInfo {
    SourceType type;
    int l_idx1, l_idx2; // Indices for Left state (e.g. closed[l_idx1][l_idx2])
    int r_idx1, r_idx2; // Indices for Right state
};

// Storage for reconstruction
// trace_merge[u][child_index][state_type][d1][d2]
// state_type: 0 = closed, 1 = open, 2 = global
// child_index: 0 is merge(Unit, Child0), 1 is merge(Res, Child1)...
// To save memory, we serialize or use a flattened vector.
// Given 128MB limit is tight-ish for heavy structs, we optimize.
// However, N=200k, we process children sequentially.
// We store trace info for every merge step.
vector<vector<vector<TraceInfo>>> trace_closed[MAXN]; 
vector<vector<TraceInfo>> trace_open[MAXN];
vector<vector<TraceInfo>> trace_global[MAXN];

// Helper to update max and record trace
void update(long long &target, long long candidate, 
            TraceInfo &trace_target, const TraceInfo &candidate_trace) {
    if (candidate > target) {
        target = candidate;
        trace_target = candidate_trace;
    }
}

// Function to shift a child's state (add distance + 1)
State shift_state(const State& s) {
    State res;
    res.global_max = s.global_max;
    
    // Shift Open
    for(int i=0; i<K; ++i) {
        if(s.open[i] != INF_LL && i + 1 < K) {
            res.open[i+1] = s.open[i];
        }
    }
    
    // Shift Closed
    for(int i=0; i<K; ++i) {
        for(int j=0; j<K; ++j) {
            if(s.closed[i][j] != INF_LL && i + 1 < K && j + 1 < K) {
                res.closed[i+1][j+1] = s.closed[i][j];
            }
        }
    }
    return res;
}

// Global path result
vector<int> final_path;

// Recursive Reconstruction
// type: 0=closed, 1=open, 2=global
// d1, d2: dimensions
void reconstruct(int u, int child_idx, int type, int d1, int d2);

// Merge Logic
State dp[MAXN];

void solve(int u, int p) {
    // Collect children
    vector<int> children;
    for(int v : adj[u]) {
        if(v != p) children.push_back(v);
    }

    // Initialize Current State with Node U unit
    State curr;
    curr.global_max = P[u];
    curr.open[0] = P[u];
    curr.closed[0][0] = P[u];

    // Allocate trace for this node
    int m = children.size();
    trace_closed[u].resize(m);
    trace_open[u].resize(m);
    trace_global[u].resize(m);

    // Initial Trace: Just unit U
    // We treat the "Unit U" as the result of "Merge Step -1".
    // But to unify, we imagine a merge loop.
    // Actually, we reconstruct from the *last* merge step backwards.
    // The "base" of recursion is when child_idx < 0, which means we hit the Unit U or Empty.
    
    // Process children
    for(int i=0; i<m; ++i) {
        int v = children[i];
        solve(v, u);
        State child = shift_state(dp[v]);
        
        State next_state;
        
        // Prepare trace storage for this step
        // Dimensions: Closed [K][K], Open [K], Global [1]
        trace_closed[u][i].resize(K * K);
        trace_open[u][i].resize(K);
        trace_global[u][i].resize(1);

        // --- MERGE LOGIC ---
        
        // 1. Inherit from Left (Curr) - simply skipping child
        // Check Closed
        for(int r=0; r<K; ++r) {
            for(int c=0; c<K; ++c) {
                if(curr.closed[r][c] != INF_LL) {
                    TraceInfo ti = {FROM_L, r, c, -1, -1};
                    update(next_state.closed[r][c], curr.closed[r][c], 
                           trace_closed[u][i][r*K + c], ti);
                }
            }
        }
        // Check Open
        for(int r=0; r<K; ++r) {
            if(curr.open[r] != INF_LL) {
                TraceInfo ti = {FROM_L, r, -1, -1, -1};
                update(next_state.open[r], curr.open[r], 
                       trace_open[u][i][r], ti);
            }
        }
        // Check Global
        if(curr.global_max != INF_LL) {
             TraceInfo ti = {FROM_L, -1, -1, -1, -1};
             update(next_state.global_max, curr.global_max, 
                    trace_global[u][i][0], ti);
        }

        // 2. Inherit from Right (Child) - replacing (only if Left was effectively empty? No, accumulative)
        // Actually, we treat 'curr' as the accumulator. 
        // We can start a new chain from 'child' and discard 'curr'?
        // Only if 'curr' was purely U and we discard U? 
        // Or if 'curr' was empty?
        // Since we want to select a subset, skipping 'curr' implies skipping U and all previous children.
        // That is covered by 'child.global_max'.
        // But can we start a closed chain from child? Yes, if we drop U.
        // Note: The problem requires path starting at 1.
        // So dropping 1 is invalid.
        // However, for subtrees, we can drop the local root U.
        // So yes, we can take Child Closed/Open as the new Closed/Open for U.
        
        // Inherit Closed from Child
        for(int r=0; r<K; ++r) {
            for(int c=0; c<K; ++c) {
                if(child.closed[r][c] != INF_LL) {
                    TraceInfo ti = {FROM_R, -1, -1, r, c};
                    update(next_state.closed[r][c], child.closed[r][c], 
                           trace_closed[u][i][r*K + c], ti);
                }
            }
        }
        // Inherit Open from Child
        for(int r=0; r<K; ++r) {
            if(child.open[r] != INF_LL) {
                TraceInfo ti = {FROM_R, -1, -1, r, -1};
                update(next_state.open[r], child.open[r], 
                       trace_open[u][i][r], ti);
            }
        }
        // Inherit Global from Child
        if(child.global_max != INF_LL) {
             TraceInfo ti = {FROM_R, -1, -1, -1, -1};
             update(next_state.global_max, child.global_max, 
                    trace_global[u][i][0], ti);
        }

        // 3. Combine Left and Right
        
        // A. Closed + Closed -> Closed
        for(int l1=0; l1<K; ++l1) {
            for(int l2=0; l2<K; ++l2) {
                if(curr.closed[l1][l2] == INF_LL) continue;
                for(int r1=0; r1<K; ++r1) {
                    for(int r2=0; r2<K; ++r2) {
                        if(child.closed[r1][r2] == INF_LL) continue;
                        
                        // L then R
                        if(l2 + 2 + r1 <= K) {
                            TraceInfo ti = {L_THEN_R, l1, l2, r1, r2};
                            update(next_state.closed[l1][r2], 
                                   curr.closed[l1][l2] + child.closed[r1][r2], 
                                   trace_closed[u][i][l1*K + r2], ti);
                        }
                        // R then L
                        if(r2 + 2 + l1 <= K) {
                             TraceInfo ti = {R_THEN_L, l1, l2, r1, r2};
                             update(next_state.closed[r1][l2], 
                                    child.closed[r1][r2] + curr.closed[l1][l2], 
                                    trace_closed[u][i][r1*K + l2], ti);
                        }
                    }
                }
            }
        }
        
        // B. Closed + Closed -> Global (Internal Loop)
        // If we connect ends of L and R, but don't expose?
        // Wait, "Closed" means exposed first and last.
        // We can join L.last to R.first. 
        // We get a chain L.first ... R.last.
        // Can we close the loop? i.e. R.last to L.first?
        // No, it's a tree. We can't cycle.
        // So Closed + Closed always yields Closed (a longer line).
        // Unless we just stop extending it? Then it becomes Global?
        // No, Global means fully contained. A Closed chain is fully contained if we decide not to extend it.
        // So any Closed state is a candidate for Global.
        
        // C. Closed + Open -> Open
        for(int l1=0; l1<K; ++l1) {
            for(int l2=0; l2<K; ++l2) {
                if(curr.closed[l1][l2] == INF_LL) continue;
                for(int r1=0; r1<K; ++r1) {
                    if(child.open[r1] == INF_LL) continue;
                    
                    // L then R (Open)
                    if(l2 + 2 + r1 <= K) {
                        TraceInfo ti = {L_THEN_R, l1, l2, r1, -1};
                        update(next_state.open[l1], 
                               curr.closed[l1][l2] + child.open[r1], 
                               trace_open[u][i][l1], ti);
                    }
                    // R (Open) then L ??? No, Open must be at end.
                }
            }
        }
        // Symmetric: R (Closed) + L (Open) -> Open
         for(int r1=0; r1<K; ++r1) {
            for(int r2=0; r2<K; ++r2) {
                if(child.closed[r1][r2] == INF_LL) continue;
                for(int l1=0; l1<K; ++l1) {
                    if(curr.open[l1] == INF_LL) continue;
                    
                    if(r2 + 2 + l1 <= K) {
                         TraceInfo ti = {R_THEN_L, l1, -1, r1, r2}; // Child first, then Curr
                         update(next_state.open[r1], 
                                child.closed[r1][r2] + curr.open[l1], 
                                trace_open[u][i][r1], ti);
                    }
                }
            }
        }
        
        // D. Open + Open -> Global (Connect two tips)
        for(int l1=0; l1<K; ++l1) {
            if(curr.open[l1] == INF_LL) continue;
            for(int r1=0; r1<K; ++r1) {
                if(child.open[r1] == INF_LL) continue;
                
                if(l1 + 2 + r1 <= K) {
                    TraceInfo ti = {L_THEN_R, l1, -1, r1, -1}; // Just merge info
                    update(next_state.global_max, 
                           curr.open[l1] + child.open[r1], 
                           trace_global[u][i][0], ti);
                }
            }
        }
        
        // E. Candidates for Global from new Closed/Open
        // Any Closed path can be considered "finished" (Global)
        // Any Open path can be considered "finished" (Global)
        // We check next_state and update next_state.global
        // But we need to record trace.
        // To avoid self-dependency, we do this after all generations?
        // Or check source components.
        // Actually simpler: At the END of merge, take max of next.global, next.open, next.closed.
        // But careful: we need to retrieve the path.
        // Let's rely on standard transitions. 
        // If "Closed" is best global, we'll pick it via a dummy transition?
        // No. "Global" state represents the "Best Internal Path".
        // A Closed path is internal if we stop extending it.
        // So we update next.global with next.closed and next.open values.
        
        for(int r=0; r<K; ++r) {
             for(int c=0; c<K; ++c) {
                 if(next_state.closed[r][c] > next_state.global_max) {
                     TraceInfo ti = {NONE, r, c, -1, -1}; // Special type?
                     // Let's use specific encoding: 
                     // Type NONE with indices r,c implies "From Closed[r][c]" of THIS state.
                     // But we are constructing THIS state.
                     // We can't reference it.
                     // We just update with same logic that created closed[r][c].
                     // But that's redundant calc.
                     // Let's just say global_max can pull from open/closed.
                     // We handle this by adding a post-check or allowing "conversion".
                 }
             }
        }
        
        // Easier way: Only strictly internal merges (Open+Open) go to Global.
        // When we finish all children, the answer is max(Global, Open, Closed).
        
        curr = next_state;
    }
    
    // Post-processing for Node U: 
    // We implicitly checked "Drop U" vs "Keep U" via FROM_R logic.
    dp[u] = curr;
}

// Function to collect nodes from a child (shifting logic reversed)
void reconstruct_child(int v, int type, int d1, int d2) {
    // We need to match the state 'child' which was shifted dp[v].
    // Shifted Closed[d1][d2] came from dp[v].closed[d1-1][d2-1]
    // Shifted Open[d1] came from dp[v].open[d1-1]
    // Global came from Global
    
    // Determine unshifted parameters
    int ud1 = d1 - 1;
    int ud2 = d2 - 1;
    
    // Determine # of children of v to find last merge index
    int m = 0;
    for(int n : adj[v]) if(n != ((adj[v].size()==1 && v!=1) ? adj[v][0] : -1)) m++; 
    // Wait, simpler: calculate m
    // Using parent pointer logic is tricky without passing parent.
    // But we traverse adj[v] same order.
    // Parent of v is unknown here.
    // But graph is tree. Parent is unique.
    // Wait, 'adj' contains parent. We need to skip it.
    // 'reconstruct' needs 'p'.
    // We'll pass 'p' to reconstruct or just handle traversal carefully.
    // Let's rebuild children list.
}

void reconstruct(int u, int child_idx, int type, int d1, int d2, int p) {
    // If child_idx < 0, we are at the base case (Node U).
    if (child_idx < 0) {
        // If we reached here, it means we used Node U.
        // The only valid base states are those involving U.
        // Unit U has Closed[0][0], Open[0], Global=P[u].
        // If type/d1/d2 matches this, we pick U.
        // If type/d1/d2 implies FROM_R logic fully replaced U, we wouldn't reach index -1.
        // (Because FROM_R consumes child and stays at index i).
        // Wait, index -1 is "Before first child". State is Unit U.
        // So we add U.
        final_path.push_back(u);
        return;
    }

    // Retrieve trace
    TraceInfo ti;
    if (type == 0) ti = trace_closed[u][child_idx][d1*K + d2];
    else if (type == 1) ti = trace_open[u][child_idx][d1];
    else ti = trace_global[u][child_idx][0];

    // Recurse based on type
    int v_idx = child_idx;
    // Helper to find actual child node v
    int v = -1, cnt = 0;
    for(int node : adj[u]) {
        if(node != p) {
            if(cnt == v_idx) { v = node; break; }
            cnt++;
        }
    }

    switch(ti.type) {
        case FROM_L:
            // Recurse Left (Previous merge step) with same params
            reconstruct(u, child_idx - 1, type, d1, d2, p);
            break;
            
        case FROM_R:
            // Recurse Right (Child v)
            // Need to unshift params for child
            if (type == 0) reconstruct(v, (int)adj[v].size() - (v==1?0:1) - 1, type, ti.r_idx1 - 1, ti.r_idx2 - 1, u);
            else if (type == 1) reconstruct(v, (int)adj[v].size() - (v==1?0:1) - 1, type, ti.r_idx1 - 1, -1, u);
            else reconstruct(v, (int)adj[v].size() - (v==1?0:1) - 1, type, -1, -1, u);
            break;
            
        case L_THEN_R:
            // Recurse Left
            if (ti.l_idx2 != -1) reconstruct(u, child_idx - 1, 0, ti.l_idx1, ti.l_idx2, p); // Closed
            else reconstruct(u, child_idx - 1, 1, ti.l_idx1, -1, p); // Open
            
            // Recurse Right
            if (ti.r_idx2 != -1) reconstruct(v, (int)adj[v].size() - (v==1?0:1) - 1, 0, ti.r_idx1 - 1, ti.r_idx2 - 1, u);
            else reconstruct(v, (int)adj[v].size() - (v==1?0:1) - 1, 1, ti.r_idx1 - 1, -1, u);
            break;
            
        case R_THEN_L:
            // Recurse Right
            if (ti.r_idx2 != -1) reconstruct(v, (int)adj[v].size() - (v==1?0:1) - 1, 0, ti.r_idx1 - 1, ti.r_idx2 - 1, u);
            else reconstruct(v, (int)adj[v].size() - (v==1?0:1) - 1, 1, ti.r_idx1 - 1, -1, u);

            // Recurse Left
            if (ti.l_idx2 != -1) reconstruct(u, child_idx - 1, 0, ti.l_idx1, ti.l_idx2, p);
            else reconstruct(u, child_idx - 1, 1, ti.l_idx1, -1, p);
            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];
    }

    solve(1, 0);

    // Find best final state
    long long ans = dp[1].global_max;
    int type = 2, d1 = -1, d2 = -1;
    
    // Check Open (valid global paths)
    for(int i=0; i<K; ++i) {
        if(dp[1].open[i] > ans) {
            ans = dp[1].open[i];
            type = 1; d1 = i; d2 = -1;
        }
    }
    // Check Closed (valid global paths)
    for(int i=0; i<K; ++i) {
        for(int j=0; j<K; ++j) {
            if(dp[1].closed[i][j] > ans) {
                ans = dp[1].closed[i][j];
                type = 0; d1 = i; d2 = j;
            }
        }
    }

    cout << ans << "\n";
    
    int child_count = adj[1].size(); // 1 is root
    reconstruct(1, child_count - 1, type, d1, d2, 0);
    
    cout << final_path.size() << "\n";
    for (int i = 0; i < final_path.size(); i++) {
        cout << final_path[i] << (i == final_path.size() - 1 ? "" : " ");
    }
    cout << "\n";

    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void solve(int, int)':
Main.cpp:154:27: error: invalid initialization of reference of type 'TraceInfo&' from expression of type '__gnu_cxx::__alloc_traits<std::allocator<std::vector<TraceInfo> >, std::vector<TraceInfo> >::value_type' {aka 'std::vector<TraceInfo>'}
  154 |                     update(next_state.closed[r][c], curr.closed[r][c],
      |                     ~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  155 |                            trace_closed[u][i][r*K + c], ti);
      |                            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:66:24: note: in passing argument 3 of 'void update(long long int&, long long int, TraceInfo&, const TraceInfo&)'
   66 |             TraceInfo &trace_target, const TraceInfo &candidate_trace) {
      |             ~~~~~~~~~~~^~~~~~~~~~~~
Main.cpp:192:27: error: invalid initialization of reference of type 'TraceInfo&' from expression of type '__gnu_cxx::__alloc_traits<std::allocator<std::vector<TraceInfo> >, std::vector<TraceInfo> >::value_type' {aka 'std::vector<TraceInfo>'}
  192 |                     update(next_state.closed[r][c], child.closed[r][c],
      |                     ~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  193 |                            trace_closed[u][i][r*K + c], ti);
      |                            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:66:24: note: in passing argument 3 of 'void update(long long int&, long long int, TraceInfo&, const TraceInfo&)'
   66 |             TraceInfo &trace_target, const TraceInfo &candidate_trace) {
      |             ~~~~~~~~~~~^~~~~~~~~~~~
Main.cpp:225:35: error: invalid initialization of reference of type 'TraceInfo&' from expression of type '__gnu_cxx::__alloc_traits<std::allocator<std::vector<TraceInfo> >, std::vector<TraceInfo> >::value_type' {aka 'std::vector<TraceInfo>'}
  225 |                             update(next_state.closed[l1][r2],
      |                             ~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~
  226 |                                    curr.closed[l1][l2] + child.closed[r1][r2],
      |                                    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  227 |                                    trace_closed[u][i][l1*K + r2], ti);
      |                                    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:66:24: note: in passing argument 3 of 'void update(long long int&, long long int, TraceInfo&, const TraceInfo&)'
   66 |             TraceInfo &trace_target, const TraceInfo &candidate_trace) {
      |             ~~~~~~~~~~~^~~~~~~~~~~~
Main.cpp:232:36: error: invalid initialization of reference of type 'TraceInfo&' from expression of type '__gnu_cxx::__alloc_traits<std::allocator<std::vector<TraceInfo> >, std::vector<TraceInfo> >::value_type' {aka 'std::vector<TraceInfo>'}
  232 |                              update(next_state.closed[r1][l2],
      |                              ~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~
  233 |                                     child.closed[r1][r2] + curr.closed[l1][l2],
      |                                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  234 |                                     trace_closed[u][i][r1*K + l2], ti);
      |                                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:66:24: note: in passing argument 3 of 'void update(long long int&, long long int, TraceInfo&, const TraceInfo&)'
   66 |             TraceInfo &trace_target, const TraceInfo &candidate_trace) {
      |             ~~~~~~~~~~~^~~~~~~~~~~~
Main.cpp: In function 'void reconstruct(int, int, int, int, int, int)':
Main.cpp:387:61: error: no match for 'operator=' (operand types are 'TraceInfo' and '__gnu_cxx::__alloc_traits<std::allocator<std::vector<TraceInfo> >, std::vector<TraceInfo> >::value_type' {aka 'std::vector<TraceInfo>'})
  387 |     if (type == 0) ti = trace_closed[u][child_idx][d1*K + d2];
      |                                                             ^
Main.cpp:46:8: note: candidate: 'constexpr TraceInfo& TraceInfo::operator=(const TraceInfo&)'
   46 | struct TraceInfo {
      |        ^~~~~~~~~
Main.cpp:46:8: note:   no known conversion for argument 1 from '__gnu_cxx::__alloc_traits<std::allocator<std::vector<TraceInfo> >, std::vector<TraceInfo> >::value_type' {aka 'std::vector<TraceInfo>'} to 'const TraceInfo&'
Main.cpp:46:8: note: candidate: 'constexpr TraceInfo& TraceInfo::operator=(TraceInfo&&)'
Main.cpp:46:8: note:   no known conversion for argument 1 from '__gnu_cxx::__alloc_traits<std::allocator<std::vector<TraceInfo> >, std::vector<TraceInfo> >::value_type' {aka 'std::vector<TraceInfo>'} to 'TraceInfo&&'