#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
using namespace std;
const int MAXN = 200005;
struct State {
int d_out;
long long profit;
};
int N, K;
vector<pair<int, int>> adj[MAXN]; // Stores {neighbor, original_index} if needed, but simple list ok
vector<int> tree[MAXN];
long long P[MAXN];
// dp[u][d_in]
vector<State> memo[MAXN][4];
bool computed[MAXN][4];
// Reconstruction helpers
struct RecStep {
int v;
int d_in_passed;
int state_idx_used;
};
map<pair<int, int>, vector<RecStep>> path_choices;
map<pair<int, int>, bool> picked_node;
void merge_states(int u, int d_in, bool pick) {
long long current_profit = (pick ? P[u] : 0);
int current_gap = (pick ? 0 : d_in);
// Track which children used. Since K is small, we just need a local used list?
// N is 200k, we can't iterate used array.
// But we only visit children that are connected.
// We can use a vector of bools aligned with tree[u].
vector<bool> child_used(tree[u].size(), false);
vector<RecStep> steps;
// Greedy loop
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];
// Access DP state
vector<State>& opts = memo[v][req_in];
for (int s_i = 0; s_i < opts.size(); s_i++) {
State& s = opts[s_i];
int next_g = s.d_out + 1; // distance adds 1 stepping back up
// Greedy Criteria: Min next gap, then Max 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 (best_child_idx != -1) {
child_used[best_child_idx] = true;
current_profit += best_gain;
current_gap = best_next_gap;
steps.push_back({tree[u][best_child_idx], req_in, best_state_idx});
} else {
break;
}
}
// Add result to DP
vector<State>& res = memo[u][d_in];
// Check if Pareto optimal
bool dominated = false;
for (auto& s : res) {
if (s.d_out <= current_gap && s.profit >= current_profit) {
dominated = true; break;
}
}
if (!dominated) {
// Remove dominated
vector<State> next_res;
next_res.push_back({current_gap, current_profit});
for (auto& s : res) {
if (current_gap <= s.d_out && current_profit >= s.profit) continue;
next_res.push_back(s);
}
res = next_res;
// If this branch was better, record the choices
// Note: This simplistic overwrite assumes the last added is best or we only add one.
// In reality, we might have multiple valid states.
// We only save the path for the specific (u, d_in) -> outcome.
// We need to map (u, d_in, resulting_d_out) -> steps?
// Simpler: Just reconstruct by re-running greedy later.
}
}
void solve(int u, int p) {
// Post-order DFS
for (int v : tree[u]) {
if (v == p) continue;
solve(v, u);
}
// Clean tree[u] to remove parent
vector<int> children;
for(int v : tree[u]) if(v != p) children.push_back(v);
tree[u] = children;
// Case 1: Pick u (Valid for any d_in)
// To simplify, we calculate "Picked" outcome once (since it resets d_in to 0).
// But we need to store it for all d_in?
// No, if picked, d_in is irrelevant for the internal logic, but relevant for the DP table key.
// The result {d_out, profit} is the same for all d_in.
// Run logic for "Picked"
merge_states(u, 0, true);
// The result is stored in memo[u][0].
// Copy this result to memo[u][1..K-1] if needed?
// Actually, distinct branches.
// Wait, if we pick u, the parent sees us returning d_out.
// The parent passed us d_in. We ignored it (reset).
// So for all d_in, we can transition to the "Picked" state.
State picked_res = memo[u][0][0]; // Assuming one best state
// Case 2: Don't pick u
for (int d = 0; d < K; d++) {
// We can only "not pick" if we continue the gap
merge_states(u, d, false);
}
// Combine "Picked" into all slots
// If Picking is better than Not Picking, use it.
for (int d = 0; d < K; d++) {
// Check if picked_res dominates existing
bool dominated = false;
for (auto& s : memo[u][d]) {
if (s.d_out <= picked_res.d_out && s.profit >= picked_res.profit) dominated = true;
}
if (!dominated) {
// Remove dominated by picked
vector<State> next_res;
next_res.push_back(picked_res);
for (auto& s : memo[u][d]) {
if (picked_res.d_out <= s.d_out && picked_res.profit >= s.profit) continue;
next_res.push_back(s);
}
memo[u][d] = next_res;
}
}
}
// Global result storage
vector<int> final_path;
void reconstruct(int u, int d_in, bool picked_parent_forced = false) {
// We need to decide if we picked u or not based on memo
// But wait, the memo merges both.
// We need to know WHICH choice led to the max profit for (u, d_in).
// Re-run the comparison to find if "Pick" or "No Pick" was better
// (Omitted for brevity, logic follows the max found in solve)
// Then call reconstruct on children using the 'steps' logic.
if (picked_u) final_path.push_back(u);
// ...
}
