/*
* Problem: Travelling Trader (CCO 2023 P2)
* Solution: Optimized Tree DP for K=1, 2, 3
* Fixes for K=2:
* - Proper handling of Incoming candidates (dp2 vs dp3).
* - Strict definition of Loop candidates (dp2 only).
* - Inclusion of Grandchild Jump for dp1.
* - O(N) transitions using Top-3.
*/
#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>
using namespace std;
const long long INF = 1e18;
int N, K;
vector<long long> P;
vector<vector<int>> adj;
// ================= K=1: Longest Path =================
long long max_profit_k1 = -1;
vector<int> path_k1;
void dfs_k1(int u, int p, long long current_profit, vector<int>& current_path) {
current_path.push_back(u);
current_profit += P[u];
if (current_profit > max_profit_k1) {
max_profit_k1 = current_profit;
path_k1 = current_path;
}
for (int v : adj[u]) {
if (v != p) {
dfs_k1(v, u, current_profit, current_path);
}
}
current_path.pop_back();
}
void solve_k1() {
vector<int> current_path;
dfs_k1(1, 0, 0, current_path);
cout << max_profit_k1 << "\n";
cout << path_k1.size() << "\n";
for (int i = 0; i < path_k1.size(); i++) cout << path_k1[i] << (i == path_k1.size()-1 ? "" : " ");
cout << "\n";
}
// ================= K=3: Euler Tour =================
void print_k3(int u, int p, bool reverse_mode) {
if (!reverse_mode) {
cout << u << " ";
for (int v : adj[u]) {
if (v != p) {
print_k3(v, u, true);
}
}
} else {
vector<int> children;
for (int v : adj[u]) if (v != p) children.push_back(v);
for (int i = children.size() - 1; i >= 0; i--) {
print_k3(children[i], u, false);
}
cout << u << " ";
}
}
void solve_k3() {
long long total = 0;
for (int i = 1; i <= N; i++) total += P[i];
cout << total << "\n";
cout << N << "\n";
print_k3(1, 0, false);
cout << "\n";
}
// ================= K=2: Complex Tree DP =================
struct DPState {
long long val;
int choice_a, choice_b, choice_c;
// a: Incoming/Loop, b: Loop, c: Tail
// Special: choice_a = -2 -> Grandchild Jump
int type_a; // 1 for dp2, 2 for dp3 (Only for choice_a)
};
DPState dp[200005][3];
struct Cand {
long long gain;
int id;
int type; // 1 for dp2, 2 for dp3
};
void update_top3(vector<Cand>& top3, Cand c) {
top3.push_back(c);
int i = top3.size() - 1;
while(i > 0) {
if(top3[i].gain > top3[i-1].gain) {
swap(top3[i], top3[i-1]);
i--;
} else break;
}
if(top3.size() > 3) top3.pop_back();
}
void dfs_k2(int u, int p) {
vector<int> children;
long long sum_p = 0;
for (int v : adj[u]) {
if (v != p) {
dfs_k2(v, u);
children.push_back(v);
sum_p += P[v];
}
}
long long base = sum_p + P[u];
vector<Cand> cands_d1, cands_d2, cands_in;
long long max_dp3_jump = -INF;
int dp3_jump_id = -1;
for (int v : children) {
// Candidates for Tail (D1)
update_top3(cands_d1, {dp[v][0].val - P[v], v, 0});
// Candidates for Loop (Strict D2)
update_top3(cands_d2, {dp[v][1].val - P[v], v, 1});
// Candidates for Incoming (Max of Strict D2, D3)
// Incoming allows ending at Child (D3) or U (D2)
long long gain_d2 = dp[v][1].val - P[v];
long long gain_d3 = dp[v][2].val - P[v];
if (gain_d3 > gain_d2) {
update_top3(cands_in, {gain_d3, v, 2});
} else {
update_top3(cands_in, {gain_d2, v, 1});
}
// Grandchild Jump (Standalone path)
if (P[u] + dp[v][2].val > max_dp3_jump) {
max_dp3_jump = P[u] + dp[v][2].val;
dp3_jump_id = v;
}
}
// --- DP1: Start u ---
dp[u][0] = {base, -1, -1, -1, 0};
// 1. Loop + Tail
for(auto& a : cands_d2) { // Loop must be Strict D2
for(auto& b : cands_d1) { // Tail is D1
if(a.id != b.id) {
if(base + a.gain + b.gain > dp[u][0].val) {
dp[u][0] = {base + a.gain + b.gain, a.id, -1, b.id, 1};
}
}
}
}
// 2. Just Loop (D2)
for(auto& a : cands_d2) {
if(base + a.gain > dp[u][0].val) dp[u][0] = {base + a.gain, a.id, -1, -1, 1};
}
// 3. Just Tail (D1)
for(auto& b : cands_d1) {
if(base + b.gain > dp[u][0].val) dp[u][0] = {base + b.gain, -1, -1, b.id, 0};
}
// 4. Grandchild Jump (Overrides all)
if (max_dp3_jump > dp[u][0].val) {
dp[u][0] = {max_dp3_jump, -2, -1, dp3_jump_id, 0};
}
// --- DP2: Start u, End u (Strict) ---
dp[u][1] = {base, -1, -1, -1, 0};
// Only Loop allowed (Tail would end deep)
for(auto& a : cands_d2) {
if(base + a.gain > dp[u][1].val) dp[u][1] = {base + a.gain, a.id, -1, -1, 1};
}
// --- DP3: Start Child ---
dp[u][2] = {-INF, -1, -1, -1, 0};
for(auto& a : cands_in) { // Incoming (D2 or D3)
long long current = base + a.gain;
// Just Incoming
if(current > dp[u][2].val) dp[u][2] = {current, a.id, -1, -1, a.type};
// Incoming + Loop
for(auto& b : cands_d2) {
if(a.id != b.id) {
if(current + b.gain > dp[u][2].val) {
dp[u][2] = {current + b.gain, a.id, b.id, -1, a.type};
}
}
}
// Incoming + Tail
for(auto& c : cands_d1) {
if(a.id != c.id) {
if(current + c.gain > dp[u][2].val) {
dp[u][2] = {current + c.gain, a.id, -1, c.id, a.type};
}
}
}
// Incoming + Loop + Tail
for(auto& b : cands_d2) {
if(a.id == b.id) continue;
for(auto& c : cands_d1) {
if(a.id != c.id && b.id != c.id) {
if(current + b.gain + c.gain > dp[u][2].val) {
dp[u][2] = {current + b.gain + c.gain, a.id, b.id, c.id, a.type};
}
}
}
}
}
}
// Reconstruction
void build_k2(int u, int p, int type, vector<int>& path);
void add_child_path(int v, int u, int type, vector<int>& path, bool reverse_output) {
vector<int> sub;
build_k2(v, u, type, sub);
if(reverse_output) {
path.insert(path.end(), sub.rbegin(), sub.rend());
} else {
path.insert(path.end(), sub.begin(), sub.end());
}
}
void build_k2(int u, int p, int type, vector<int>& path) {
DPState& s = dp[u][type];
// Grandchild Jump
if (type == 0 && s.choice_a == -2) {
path.push_back(u);
add_child_path(s.choice_c, u, 2, path, false); // Forward DP3
return;
}
int va = s.choice_a;
int vb = s.choice_b;
int vc = s.choice_c;
auto add_singles = [&]() {
for (int v : adj[u]) {
if (v != p && v != va && v != vb && v != vc) {
path.push_back(v);
}
}
};
if (type == 0) { // DP1: u -> Rev(va, D2) -> Singles -> vc(D1)
path.push_back(u);
if (va != -1) add_child_path(va, u, 1, path, true);
add_singles();
if (vc != -1) add_child_path(vc, u, 0, path, false);
}
else if (type == 1) { // DP2: u -> Rev(va, D2) -> Singles
path.push_back(u);
if (va != -1) add_child_path(va, u, 1, path, true);
add_singles();
}
else if (type == 2) { // DP3: va(In) -> u -> Rev(vb, D2) -> Singles -> vc(D1)
if (va != -1) {
// Incoming can be D2 or D3
if (s.type_a == 1) add_child_path(va, u, 1, path, false); // D2 Forward
else add_child_path(va, u, 2, path, true); // Rev(Rev(D3)) = D3 Reversed?
// Wait. Incoming is Rev(Rev(path))?
// Incoming structure: Path -> u.
// D2(Forward) is u->child. Rev(D2) is child->u. This works.
// D3(Forward) is child->child. Rev(D3) is child->child.
// But we need to END at u.
// Dist(End(Rev(D3)), u) = Dist(child, u) = 2. Valid.
// So we need Rev(D3).
// D3 standard output is Forward. So we need Reverse output.
}
path.push_back(u);
if (vb != -1) add_child_path(vb, u, 1, path, true); // Loop Rev(D2)
add_singles();
if (vc != -1) add_child_path(vc, u, 0, path, false); // Tail D1
}
}
void solve_k2() {
dfs_k2(1, 0);
cout << dp[1][0].val << "\n";
vector<int> path;
build_k2(1, 0, 0, path);
cout << path.size() << "\n";
for(size_t i=0; i<path.size(); ++i) cout << path[i] << (i==path.size()-1?"":" ");
cout << "\n";
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL);
if (!(cin >> N >> K)) return 0;
P.resize(N + 1);
adj.resize(N + 1);
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];
if (K == 1) solve_k1();
else if (K == 2) solve_k2();
else solve_k3();
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... |