#include "highway.h"
#include <queue>
#include <algorithm>
#define int long long
#define fi first
#define se second
#define ll int
using namespace std;
// progression
// solve sub2 (done!), sub4 (done!) then try sub6
vector<int> dist, parEdge, col, colVertex, heavy;
vector<vector<pair<int, int>>> adj, adj_t;
void dfs(int u, int d, int c, int par = -1){
dist[u] = d;
for (pair<int, int> pa: adj[u]){
int v = pa.fi;
if (v == par) continue;
parEdge[v] = pa.se;
dfs(v, d+1, c, u);
}
}
vector<int> u_list, v_list;
int sub2(int n, int m, int a, int b, int inSub, int r_edge) {
vector<int> cand;
for (int i = 0; i < n; i++) {
if (colVertex[i] == inSub) cand.push_back(i);
}
// Sort by distance descending (furthest nodes first)
sort(cand.begin(), cand.end(), [](int i, int j) {
return dist[i] > dist[j];
});
// Calculate a stable newBase where the target node is 'collapsed' to the root
vector<signed> w(m, 1);
for (int i = 0; i < m; i++) {
int u_i = u_list[i], v_i = v_list[i];
// Zero out EVERYTHING in the other subtree and EVERYTHING in this subtree
if (colVertex[u_i] != inSub && colVertex[v_i] != inSub && i != r_edge) w[i] = 0;
if (colVertex[u_i] == inSub && colVertex[v_i] == inSub) w[i] = 0;
// Basically: only r_edge and cross-edges are 1
}
// Simplest approach: Just zero out all tree edges in the whole graph except r_edge
fill(w.begin(), w.end(), 1);
for(int i=0; i<n; i++) if(i != u_list[r_edge] && i != v_list[r_edge]) w[parEdge[i]] = 0;
int newBase = ask(w);
int l = 0, r = (int)cand.size() - 2; // root is always a candidate but dist 0
int res = cand.back(); // Default to root
while (l <= r) {
int mid = (l + r) / 2;
fill(w.begin(), w.end(), 1);
// 1. Collapsing the ENTIRE other subtree
for (int i = 0; i < n; i++) {
if (colVertex[i] != inSub && i != u_list[r_edge] && i != v_list[r_edge])
w[parEdge[i]] = 0;
}
// 2. Collapsing only the candidates from mid+1 to end in current subtree
for (int i = mid + 1; i < cand.size(); i++) {
if (cand[i] != u_list[r_edge] && cand[i] != v_list[r_edge])
w[parEdge[cand[i]]] = 0;
}
if (ask(w) > newBase) {
res = cand[mid];
r = mid - 1;
} else {
l = mid + 1;
}
}
return res;
}
void find_pair(signed n, vector<signed> u, vector<signed> v, signed ap, signed bp) {
for (int ui: u){
u_list.push_back(ui);
}
for (int vi: v){
v_list.push_back(vi);
}
int a = (int)ap, b = (int) bp;
int m = u.size();
vector<signed> w(m, 0);
int base = ask(w);
int totLen = base / (int)a;
col.assign(m, 0);
colVertex.assign(n, 0);
dist.assign(n, 0);
parEdge.assign(n, 0);
heavy.assign(m, 1);
adj.assign(n, vector<pair<int, int>>());
adj_t.assign(n, vector<pair<int, int>>());
// binary search for the edge on
int l = 0, r = m-1;
while(l < r){
int mid = (l + r)/2;
// for l to mid, check if increase or not
fill(w.begin(), w.end(), 0);
for (int i = l; i <= mid; i++){
w[i] = 1;
}
int chck = ask(w);
if (chck > totLen*a){
// move to r
r = mid;
// cout << "move r to " << mid << endl;
}
else{
l = mid+1;
// cout << "move l to " << mid+1 << endl;
}
}
for (int i = 0; i < m; i++){
adj_t[u[i]].push_back({v[i], i});
adj_t[v[i]].push_back({u[i], i});
}
// bfs-tree building ()
queue<pair<int, pair<int, int>>> q;
vector<int> vis(n, 0);
heavy[r] = 0;
q.push({u[r], {1, 0}});
q.push({v[r], {2, 0}});
vis[u[r]] = 1;
vis[v[r]] = 1;
while (!q.empty()){
pair<int, pair<int, int>> paun = q.front();
int un = paun.fi;
int c = paun.se.fi, d = paun.se.se;
colVertex[un] = c;
q.pop();
for (pair<int, int> pa: adj_t[un]){
int vn = pa.fi;
if (vis[vn]) continue;
vis[vn] = 1;
heavy[pa.se] = 0;
dist[vn] = d+1;
parEdge[vn] = pa.se;
// insert edge from u->v
col[pa.se] = c;
adj[un].push_back({vn, pa.se});
adj[vn].push_back({un, pa.se});
q.push({vn, {c, d+1}});
}
}
int s = sub2(n, m, a, b, 1, r);
int t = sub2(n, m, a, b, 2, r);
answer(s, t);
}
