#include "closing.h"
#ifdef LOCAL
#include "Debug.h"
#else
#define debug(...) 42
#endif
#include <bits/stdc++.h>
using namespace std;
const int N = 2e5 + 5;
const long long oo = (1LL << 62) - 10;
int n;
vector<pair<int, int>> a[N];
vector<int> idPath, idTree, trees[N];
vector<long long> bfs(int s)
{
vector<long long> dist(n, -1);
queue<int> q;
dist[s] = 0;
q.push(s);
while (!empty(q))
{
int x = q.front();
q.pop();
for (auto [y, w] : a[x])
if (dist[y] < 0)
{
dist[y] = dist[x] + w;
q.push(y);
}
}
return dist;
}
int findPath(int x, int goal, int par, vector<int> &path)
{
path.push_back(x);
if (x == goal)
return 1;
for (auto [y, _] : a[x])
if (y != par)
{
if (findPath(y, goal, x, path))
return 1;
}
path.pop_back();
return 0;
}
void dfs(int x)
{
trees[idTree[x]].push_back(x);
for (auto [y, _] : a[x])
if (idTree[y] < 0)
{
idTree[y] = idTree[x];
dfs(y);
}
}
int max_score(int N, int A, int B, long long budget, vector<int> U, vector<int> V, vector<int> W)
{
n = N;
for (int i = 0; i < n; i++)
a[i].clear();
for (int i = 0; i < size(U); i++)
{
a[U[i]].push_back({V[i], W[i]});
a[V[i]].push_back({U[i], W[i]});
}
auto distA = bfs(A);
auto distB = bfs(B);
auto distAB = distA[B];
if (distAB > budget * 2)
{
vector<long long> allDists = distA;
for (auto d : distB)
allDists.push_back(d);
sort(begin(allDists), end(allDists));
int ans = 0;
for (auto d : allDists)
if (d <= budget)
{
ans++;
budget -= d;
}
return ans;
}
// O(n^3)
vector<int> path;
findPath(A, B, -1, path);
idPath.assign(n, -1);
idTree.assign(n, -1);
int len = size(path);
for (int i = 0; i < len; i++)
{
idPath[path[i]] = idTree[path[i]] = i;
trees[i].clear();
}
for (int x : path)
dfs(x);
int ans = 0;
for (int l = 0; l < len; l++)
{
vector<long long> f(n * 2 + 1, oo);
f[0] = 0;
int cnt = 0;
for (int r = len - 1; r >= 0; r--)
{
long long curCost = 0;
int curAns = 0;
vector<long long> c1;
for (int i = 0; i < n; i++)
{
if (idTree[i] > l && idTree[i] < r)
continue;
if (idTree[i] < r) // A
{
if (idPath[i] < 0) c1.push_back(distA[i]);
else curCost += distA[i], curAns++;
}
else if (idTree[i] > l) // B
{
if (idPath[i] < 0) c1.push_back(distB[i]);
else curCost += distB[i], curAns++;
}
else // AB
{
if (idPath[i] >= 0)
curCost += max(distA[i], distB[i]), curAns += 2;
}
}
if (curCost > budget)
break;
long long rem = budget - curCost;
// add new tree to f
if (r <= l)
{
for (int x : trees[r])
if (idPath[x] < 0)
{
auto minDist = min(distA[x], distB[x]);
auto maxDist = max(distA[x], distB[x]);
for (int i = cnt; i >= 0; i--)
{
f[i + 1] = min(f[i + 1], f[i] + minDist);
f[i + 2] = min(f[i + 2], f[i] + maxDist);
}
cnt += 2;
}
}
sort(begin(c1), end(c1));
vector<long long> costFor(size(c1) + 1);
for (int i = 0; i < size(c1); i++)
costFor[i + 1] = costFor[i] + c1[i];
for (int i = 0; i <= cnt; i++)
if (f[i] > rem) break;
else
{
int u = upper_bound(begin(costFor), end(costFor), rem - f[i]) - begin(costFor);
ans = max(ans, curAns + i + u - 1);
}
}
}
return ans;
}
