#include "closing.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using Graph = vector<vector<pair<int,ll>>>;
void dfs(int u, int fu, const Graph& g, vector<ll>& dists) {
for (auto [v, w] : g[u]) {
if (v == fu) continue;
assert(dists[v] < 0);
dists[v] = dists[u] + w;
dfs(v, u, g, dists);
}
}
vector<ll> get_dists(int u, const Graph& g) {
int n = g.size();
vector<ll> dists(n, -1);
dists[u] = 0;
dfs(u, -1, g, dists);
return dists;
}
int no_common(int n, ll k, const vector<ll>& dx, const vector<ll>& dy) {
vector<ll> all;
for (auto d : dx) all.push_back(d);
for (auto d : dy) all.push_back(d);
sort(all.begin(), all.end());
int cnt = 0;
for (auto d : all) {
if (d <= k) {
k -= d;
++cnt;
} else break;
}
return cnt;
}
// linear graph
int sub4(int n, ll k, int x, int y,
const vector<ll>& dx, const vector<ll>& dy) { // x < y
int res = no_common(n, k, dx, dy);
// 0 .. [xl .. x .. xr]
// [yl .. y .. yr] .. n-1
// xr >= yl
for (int yl = x; yl <= y; ++yl) {
ll can = k;
int cur = 0;
// 0. for both x and y: visit yl
can -= max(dx[yl], dy[yl]);
cur += 2;
// 1. for x:
// 1.1. visit all of [x, yl)
for (int i = x; i < yl; ++i) {
can -= dx[i];
cur++;
}
// 1.2. option to visit others
vector<ll> rem_dx;
for (int i = 0; i < x; ++i) {
rem_dx.push_back(dx[i]);
}
for (int i = yl + 1; i < n; ++i) {
rem_dx.push_back(max(0ll, dx[i] - dy[i]));
}
// 2. for y:
// 2.1. visit all of (yl, y]
for (int i = yl+1; i <= y; i++) {
can -= dy[i];
cur++;
}
// 2.2. option to visit others
vector<ll> rem_dy;
for (int i = y+1; i < n; ++i) {
rem_dy.push_back(dy[i]);
}
if (yl == x) {
for (int i = 0; i < x; ++i) {
rem_dy.push_back(max(dy[i] - dx[i], 0ll));
}
}
if (can >= 0) {
res = max(res, no_common(n, can, rem_dx, rem_dy) + cur);
if (res == 98) {
DEBUG(x);
DEBUG(y);
DEBUG(yl);
DEBUG(cur);
throw 1;
}
}
}
return res;
}
int max_score(int n, int X, int Y, long long K,
vector<int> U, vector<int> V, vector<int> W) {
Graph g(n);
assert(int(U.size()) == n-1);
assert(int(V.size()) == n-1);
assert(int(W.size()) == n-1);
for (int i = 0; i < n - 1; i++) {
int u = U[i];
int v = V[i];
int w = W[i];
g[u].emplace_back(v, w);
g[v].emplace_back(u, w);
}
if (X > Y) swap(X, Y);
auto distsX = get_dists(X, g);
auto distsY = get_dists(Y, g);
return no_common(n, K, distsX, distsY);
}
Compilation message
closing.cpp: In function 'int sub4(int, ll, int, int, const std::vector<long long int>&, const std::vector<long long int>&)':
closing.cpp:92:17: error: 'DEBUG' was not declared in this scope
92 | DEBUG(x);
| ^~~~~
