#include "highway.h"
#include <vector>
#include <queue>
#include <algorithm>
#include <numeric>
using namespace std;
namespace {
int N, M;
vector<int> U, V;
vector<vector<int>> adj;
vector<int> w;
long long D0;
vector<int> bfs(int src) {
vector<int> dist(N, -1);
queue<int> q;
dist[src] = 0;
q.push(src);
while (!q.empty()) {
int u = q.front();
q.pop();
for (int v : adj[u]) {
if (dist[v] == -1) {
dist[v] = dist[u] + 1;
q.push(v);
}
}
}
return dist;
}
bool has_edge_prefix(int k) {
// edges [0..k] light, others heavy
for (int i = 0; i < M; ++i) w[i] = (i <= k ? 0 : 1);
return ask(w) == D0;
}
bool has_vertex_prefix(const vector<int>& pos, int k) {
// vertices with pos < k are "inside"; edges fully inside are light
for (int i = 0; i < M; ++i) {
w[i] = (pos[U[i]] < k && pos[V[i]] < k) ? 0 : 1;
}
return ask(w) == D0;
}
int endpoint_from_root(int root) {
vector<int> dist = bfs(root);
vector<int> ord(N);
iota(ord.begin(), ord.end(), 0);
sort(ord.begin(), ord.end(), [&](int a, int b) {
if (dist[a] != dist[b]) return dist[a] < dist[b];
return a < b;
});
vector<int> pos(N);
for (int i = 0; i < N; ++i) pos[ord[i]] = i;
int lo = 0; // false
int hi = N; // true
while (hi - lo > 1) {
int mid = lo + (hi - lo) / 2; // prefix size
if (has_vertex_prefix(pos, mid)) hi = mid;
else lo = mid;
}
return ord[hi - 1];
}
}
void find_pair(int n, vector<int> u, vector<int> v, int A, int B) {
(void)A; (void)B; // not needed by the algorithm beyond the D0 equality test
N = n;
U = move(u);
V = move(v);
M = (int)U.size();
adj.assign(N, {});
for (int i = 0; i < M; ++i) {
adj[U[i]].push_back(V[i]);
adj[V[i]].push_back(U[i]);
}
w.assign(M, 0);
D0 = ask(w); // all light
// Find an edge on some shortest-hop S-T path via monotone edge-prefix search.
int lo = -1; // (conceptually) none light -> false
int hi = M - 1; // all light -> true
while (hi - lo > 1) {
int mid = lo + (hi - lo) / 2;
if (has_edge_prefix(mid)) hi = mid;
else lo = mid;
}
int e = hi;
int root = U[e];
int s = endpoint_from_root(root);
int t = endpoint_from_root(s);
answer(s, t);
}
