#include "closing.h"
#include <vector>
#include <utility>
#include <algorithm>
#include <limits>
#include <queue>
using namespace std;
int max_score(int N, int X, int Y, long long K,
std::vector<int> U, std::vector<int> V, std::vector<int> W)
{
int ans = 0;
vector<vector<pair<int, int>>> adj(N);
vector<long long> distX(N), distY(N);
for (int i = 0; i < N - 1; i++)
{
adj[U[i]].emplace_back(V[i], W[i]);
adj[V[i]].emplace_back(U[i], W[i]);
}
{
auto dfs = [&](auto self, int u, vector<long long> &d, vector<bool> &vis) -> void
{
vis[u] = true;
for (auto [v, w] : adj[u])
if (!vis[v])
{
d[v] = d[u] + w;
self(self, v, d, vis);
}
};
vector<bool> vis(N, false);
distX[X] = 0;
dfs(dfs, X, distX, vis);
distY[Y] = 0;
fill(vis.begin(), vis.end(), false);
dfs(dfs, Y, distY, vis);
}
// Case 1: no intersections
{
long long remain = K;
int score = 0;
vector<long long> single;
single.insert(single.end(), distX.begin(), distX.end());
single.insert(single.end(), distY.begin(), distY.end());
sort(single.begin(), single.end());
for (auto cost : single)
if (remain >= cost)
score++, remain -= cost;
ans = max(ans, score);
}
// Case 2: intersecting
{
long long remain = K;
int root = 0;
int score = 0;
{
long long min_dist = numeric_limits<long long>::max();
for (int i = 0; i < N; i++)
if (min_dist > max(distX[i], distY[i]))
min_dist = max(distX[i], distY[i]), root = i;
vector<bool> take_two(N, false);
vector<pair<long long, int>> one, two;
priority_queue<long long> used_two;
for (int i = 0; i < N; i++)
if (i == root)
score += 2, remain -= min_dist;
else
{
long long first = min(distX[i], distY[i]);
long long second = max(distX[i], distY[i]) - first;
if (distX[i] + distY[i] == distX[root] + distY[root])
score++, remain -= first, one.emplace_back(second, i);
else
{
if (second >= first)
one.emplace_back(first, i), one.emplace_back(second, i);
else
one.emplace_back(first, i), two.emplace_back(first + second, i);
}
}
sort(one.begin(), one.end(), greater<>());
sort(two.begin(), two.end(), greater<>());
while (remain >= 0)
{
for (size_t t = 1; t <= 2; t++)
while (one.size() >= t && take_two[(one.end() - t)->second])
one.erase(one.end() - t);
// Case 1: as-is
ans = max(ans, score);
// Case 1: take one
if (!one.empty() && remain - one.back().first >= 0)
ans = max(ans, score + 1);
// Case 2: two-bundle decomposing and it is one of the previous added bundle
if (!two.empty())
{
auto [cost, i] = two.back();
long long revert_earn = max(distX[i], distY[i]) - min(distX[i], distY[i]);
if (!used_two.empty())
revert_earn = max(revert_earn, used_two.top());
if (remain - cost + revert_earn >= 0)
ans = max(ans, score + 1);
}
// Advance greedy solution
if (!two.empty())
{
auto [cost, i] = two.back();
if (one.size() < 2 || one.back().first + (one.rbegin() + 1)->first >= cost)
{
score += 2;
take_two[i] = true;
remain -= cost;
two.pop_back();
used_two.emplace(max(distX[i], distY[i]) - min(distX[i], distY[i]));
continue;
}
}
if (one.empty())
break;
auto [cost, i] = one.back();
score++;
remain -= cost;
one.pop_back();
}
}
}
return ans;
}
