#include "bits/extc++.h"
using namespace std;
template <typename T, typename... U>
void dbgh(const T& t, const U&... u) {
cerr << t;
((cerr << " | " << u), ...);
cerr << endl;
}
#ifdef DEBUG
#define dbg(...) \
cerr << "L" << __LINE__ << " [" << #__VA_ARGS__ << "]" \
<< ": "; \
dbgh(__VA_ARGS__)
#else
#define cerr \
if (false) \
cerr
#define dbg(...)
#endif
using ll = long long;
#define endl "\n"
#define long int64_t
#define sz(x) int(std::size(x))
template <typename T>
ostream& operator<<(ostream& out, const vector<T>& arr) {
out << "[";
for (int i = 0; i < sz(arr); i++) {
if (i) {
out << ", ";
}
out << arr[i];
}
return out << "]";
}
template <typename Cb>
struct Cmp {
Cb cb;
Cmp(Cb cb) : cb(cb) {}
template <typename T>
bool operator()(const T& a, const T& b) const {
return cb(a) < cb(b);
}
};
vector<vector<pair<int, long>>> edges_to_adj(
int n,
const vector<tuple<int, int, long>>& edges) {
vector<vector<pair<int, long>>> graph(n);
for (auto& [u, v, w] : edges) {
graph[u].emplace_back(v, w);
graph[v].emplace_back(u, w);
}
return graph;
}
struct DistDFS {
vector<long> dist;
vector<vector<pair<int, long>>> graph;
DistDFS(int root, int n, const vector<tuple<int, int, long>>& edges)
: dist(n), graph(edges_to_adj(n, edges)) {
dfs(root, -1, 0);
}
void dfs(int u, int p, long d) {
dist[u] = d;
for (auto& [v, w] : graph[u]) {
if (v == p) {
continue;
}
dfs(v, u, d + w);
}
}
};
struct PathDFS {
int n;
vector<int> path;
vector<char> on_path;
vector<vector<int>> path_subs;
vector<vector<pair<int, long>>> graph;
PathDFS(int u0, int u1, int n, const vector<tuple<int, int, long>>& edges)
: n(n), on_path(n), graph(edges_to_adj(n, edges)) {
pdfs(u0, -1, u1);
for (auto& a : path) {
on_path[a] = true;
}
for (auto& a : path) {
path_subs.emplace_back();
dfs(a, -1, path_subs.back());
}
}
vector<int> st;
void pdfs(int u, int p, int targ) {
st.push_back(u);
if (u == targ) {
path = st;
}
for (auto& [v, _w] : graph[u]) {
if (v == p) {
continue;
}
pdfs(v, u, targ);
}
st.pop_back();
}
void dfs(int u, int p, vector<int>& out) {
out.push_back(u);
for (auto& [v, _w] : graph[u]) {
if (v == p || on_path[v]) {
continue;
}
dfs(v, u, out);
}
}
};
struct MArr {
vector<long> vals;
vector<int> comp;
MArr(const vector<long>& vals) : vals(vals), comp(sz(vals)) {
vector<pair<long, int>> v;
for (int i = 0; i < sz(vals); i++) {
v.emplace_back(vals[i], i);
}
sort(begin(v), end(v));
for (int i = 0; i < sz(v); i++) {
comp[v[i].second] = i;
}
}
};
struct Node {
long sum, last0, last1;
int cnt;
Node operator+(const Node& n) const {
if (n.last1 == -1) {
return {sum + n.sum, last0, last1, cnt + n.cnt};
} else if (n.last0 == -1) {
return {sum + n.sum, last1, n.last1, cnt + n.cnt};
} else {
return {sum + n.sum, n.last0, n.last1, cnt + n.cnt};
}
}
static Node c_from(long x) {
return {x, -1, x, 1};
}
static Node c_def() {
return {0, -1, -1, 0};
}
};
struct ST {
int n;
vector<Node> v;
ST(int _n) : n(1 << (__lg(_n) + 1)), v(2 * n, Node::c_def()) {}
void update(int ind, const Node& x) {
ind += n;
v[ind] = x;
ind >>= 1;
for (; ind; ind >>= 1) {
v[ind] = v[ind * 2] + v[ind * 2 + 1];
}
}
Node query_pref(int ind) const {
if (ind < 0) {
return Node::c_def();
} else if (ind >= n - 1) {
return v[1];
}
Node ans = Node::c_def();
ind += n + 1;
for (; ind > 1; ind >>= 1) {
if (ind & 1) {
ans = v[ind ^ 1] + ans;
}
}
}
Node query_all() const {
return v[1];
}
template <typename Cb>
pair<int, Node> bsearch(int o, int l, int r, const Node& pref, const Cb& cb)
const {
if (l == r) {
return {l - 1, pref};
}
int mid = (l + r) / 2, lc = o * 2, rc = lc + 1;
if (!cb(pref + v[lc])) {
return bsearch(lc, l, mid, pref, cb);
} else {
return bsearch(rc, mid + 1, r, pref + v[lc], cb);
}
}
template <typename Cb>
pair<int, Node> bsearch(const Cb& cb) const {
if (cb(v[1])) {
return {n, v[1]};
}
return bsearch(1, 0, n - 1, Node::c_def(), cb);
}
};
struct DS {
MArr arr[2];
ST v_st[2];
set<int> v_inds[2];
DS(const vector<long>& v0, const vector<long>& v1)
: arr {v0, v1}, v_st {sz(v0), sz(v1)} {}
void insert(int ind, int x) {
dbg("+", ind, arr[ind].vals[x]);
v_inds[ind].insert(arr[ind].comp[x]);
v_st[ind].update(arr[ind].comp[x], Node::c_from(arr[ind].vals[x]));
}
void erase(int ind, int x) {
dbg("-", ind, arr[ind].vals[x]);
v_inds[ind].insert(arr[ind].comp[x]);
v_st[ind].update(arr[ind].comp[x], Node::c_def());
}
int query(long kv) {
int ans = -1e9;
auto upd_ans = [&](int ind) -> void {
ind = clamp(ind, -1, v_st[0].n);
auto q0 = v_st[0].query_pref(ind),
q1 = v_st[1]
.bsearch([&](const Node& o) -> bool {
return q0.sum + o.sum <= kv;
})
.second;
if (q0.sum + q1.sum <= kv) {
ans = max(ans, q0.cnt * 2 + q1.cnt);
}
};
int l = -1, r = v_st[0].n;
while (r - l > 1) {
int mid = (l + r) / 2;
auto q0 = v_st[0].query_pref(mid),
q1 = v_st[1]
.bsearch([&](const Node& o) -> bool {
return q0.sum + o.sum <= kv;
})
.second;
if (q0.sum + q1.sum <= kv &&
((q1.last0 == -1 && q1.last1 * 2 > q0.last1) ||
(q1.last0 != -1 && q1.last0 + q1.last1 > q0.last1))) {
l = mid;
} else {
r = mid;
}
}
auto upd_ans2 = [&](int ind) -> void {
upd_ans(ind - 1);
upd_ans(ind);
upd_ans(ind + 1);
};
upd_ans(l);
// upd_ans(r);
upd_ans(-1);
{
int u =
v_st[0]
.bsearch([&](const Node& o) -> bool { return o.sum <= kv; })
.first;
upd_ans(u);
upd_ans(u - 1);
}
// upd_ans2(
// v_st[0]
// .bsearch([&](const Node& o) -> bool { return o.sum <= kv; })
// .first);
{
int u = v_st[0]
.bsearch([&](const Node& o) -> bool {
return o.sum <= kv - v_st[1].query_all().sum;
})
.first;
upd_ans(u);
upd_ans(u + 1);
}
// upd_ans2(v_st[0]
// .bsearch([&](const Node& o) -> bool {
// return o.sum <= kv - v_st[1].query_all().sum;
// })
// .first);
auto upd_rep = [&](auto cb) -> void {
int u = l;
for (int it = 0; it < 2; it++) {
auto n_u = cb(u);
if (!n_u) {
return;
}
u = n_u.value();
upd_ans(u);
}
};
upd_rep([&](int u) -> optional<int> {
auto it = v_inds[0].lower_bound(u);
if (it == v_inds[0].begin()) {
return {};
}
return *(--it);
});
upd_rep([&](int u) -> optional<int> {
auto it = v_inds[0].upper_bound(u);
if (it == v_inds[0].end()) {
return {};
}
return *it;
});
return ans;
}
};
int solve_disjoint(long kv,
const vector<long>& dist0,
const vector<long>& dist1) {
vector<long> dists;
dists.insert(dists.end(), begin(dist0), end(dist0));
dists.insert(dists.end(), begin(dist1), end(dist1));
sort(begin(dists), end(dists));
long sum = 0;
int i;
for (i = 0; i < sz(dists); i++) {
sum += dists[i];
if (sum > kv) {
break;
}
}
return i;
}
int solve(int n, int u0, int u1, long kv, vector<tuple<int, int, long>> edges) {
auto dist0 = DistDFS(u0, n, edges).dist, dist1 = DistDFS(u1, n, edges).dist;
auto path_dfs = PathDFS(u0, u1, n, edges);
auto path = path_dfs.path_subs;
int m = sz(path);
vector<long> cost1(n), cost2(n), delta(n);
for (int i = 0; i < n; i++) {
cost1[i] = min(dist0[i], dist1[i]);
cost2[i] = max(dist0[i], dist1[i]);
delta[i] = cost2[i] - cost1[i];
}
map<long, vector<int>> mp;
for (int i = 0; i < m; i++) {
mp[delta[path[i][0]]].push_back(i);
}
DS ds(cost2, cost1);
int ans = solve_disjoint(kv, dist0, dist1);
dbg(ans);
int ans_add = 0;
long min_cost = 0;
for (auto& a : path_dfs.path) {
ans_add++;
min_cost += cost1[a];
}
auto move_vals = [&](vector<int> nodes, bool undo) -> void {
sort(begin(nodes), end(nodes),
Cmp([&](int u) -> long { return cost1[u]; }));
for (auto& u : nodes) {
if (path_dfs.on_path[u]) {
ans_add++;
min_cost += delta[u];
} else {
ds.erase(1, u);
ds.insert(0, u);
}
ans = max(ans, ans_add + ds.query(kv - min_cost));
dbg(ans, ans_add, min_cost);
}
if (!undo) {
return;
}
for (auto& u : nodes) {
if (path_dfs.on_path[u]) {
ans_add--;
min_cost -= delta[u];
} else {
ds.insert(1, u);
ds.erase(0, u);
}
}
};
for (int i = 0; i < n; i++) {
if (path_dfs.on_path[i]) {
continue;
}
ds.insert(1, i);
}
for (auto& [_k, vals] : mp) {
dbg(vals);
if (sz(vals) == 1) {
move_vals(path[vals[0]], false);
continue;
}
assert(sz(vals) == 2);
dbg(ans_add, min_cost);
move_vals(path[vals[0]], true);
dbg(ans_add, min_cost);
move_vals(path[vals[1]], false);
move_vals(path[vals[0]], false);
}
return ans;
}
int max_score(int n,
int u0,
int u1,
ll kv,
vector<int> edges_u,
vector<int> edges_v,
vector<int> edges_w) {
{
DS ds {{6}, {5}};
ds.insert(1, 0);
ds.erase(1, 0);
ds.insert(0, 0);
ds.insert(1, 0);
ds.erase(0, 0);
ds.erase(1, 0);
ds.insert(0, 0);
dbg(ds.query(6));
// return -1;
}
vector<tuple<int, int, long>> edges;
for (int i = 0; i < n - 1; i++) {
edges.emplace_back(edges_u[i], edges_v[i], edges_w[i]);
}
return solve(n, u0, u1, kv, edges);
}
Compilation message
closing.cpp: In member function 'int DS::query(int64_t)':
closing.cpp:301:14: warning: variable 'upd_ans2' set but not used [-Wunused-but-set-variable]
301 | auto upd_ans2 = [&](int ind) -> void {
| ^~~~~~~~
closing.cpp: In member function 'Node ST::query_pref(int) const':
closing.cpp:213:5: warning: control reaches end of non-void function [-Wreturn-type]
213 | }
| ^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1067 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1044 ms |
98544 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1077 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1077 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1077 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1067 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1067 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1067 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1067 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1067 ms |
348 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
