IOI '18 - Highway (51pts)

#include "highway.h"
#include <iostream>
#include <algorithm>
#include <cassert>

using namespace std;

vector<vector<pair<int, int>>> g;
vector<int> parent, edgeUp, depth;

void rootDFS(int node, int d = 0, int eUP = -1, int par = -1) {
  parent[node] = par;
  edgeUp[node] = eUP;
  depth[node] = d;
  for (auto [child, e] : g[node]) {
    if (child != par) rootDFS(child, d+1, e, node);
  }
}
void rootTree(int root, int N) {
  parent.resize(N);
  edgeUp.resize(N);
  depth.resize(N);
  rootDFS(root);
}
vector<int> nodesDepth(int d) {
  vector<int> v;
  for (int i = 0; i < (int)depth.size(); ++i) {
    if (depth[i] == d) v.push_back(i);
  }
  return v;
}

void find_pair(int N, vector<int> U, vector<int> V, int A, int B) {
  g.resize(N);
  int M = U.size();
  for (int i = 0; i < M; ++i) {
    g[U[i]].emplace_back(V[i], i);
    g[V[i]].emplace_back(U[i], i);
  }
  vector<int> w(M);
  long long toll = ask(w);

  rootTree(0, N);
  // Find the depth of the lower of the two nodes
  int l = 0, r = *max_element(depth.begin(), depth.end()); // l < d <= r
  while (l+1 < r) {
    int mid = (l+r)/2;
    // set all nodes lower than depth mid to have B edges leading up
    for (int i = 0; i < N; ++i) {
      if (depth[i] > mid) w[edgeUp[i]] = 1;
    }
    long long res = ask(w);
    for (int i = 0; i < N; ++i) {
      if (depth[i] > mid) w[edgeUp[i]] = 0;
    }
    if (res == toll) r = mid;
    else l = mid;
  }
  int lowerDepth = r;

  auto findWithDepth = [&toll, &w](int depth) {
    vector<int> candidates = nodesDepth(depth);
    int l = 0, r = (int)candidates.size(); // in range [l, r)
    while (l+1 < r) {
      int mid = (l+r)/2;
      for (int i = mid; i < r; ++i) w[edgeUp[candidates[i]]] = 1, assert(edgeUp[candidates[i]] >= 0);
      long long res = ask(w);
      for (int i = mid; i < r; ++i) w[edgeUp[candidates[i]]] = 0;
      if (res == toll) r = mid;
      else l = mid;
    }
    return candidates[l];
  };

  // Find the lower of the two nodes (or one of them if they have the same depth)
  int S = findWithDepth(lowerDepth);

  // Find the other node by rooting the tree at the first node
  rootTree(S, N);
  assert(toll%A==0);
  int T = findWithDepth(toll/A);
  assert(depth[T] == toll/A);
  answer(S, T);
}