#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 time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
