#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); }