Advertisement
erek1e

IOI '18 - Highway (51pts)

Jul 27th, 2022
592
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.27 KB | None | 0 0
  1. #include "highway.h"
  2. #include <iostream>
  3. #include <algorithm>
  4. #include <cassert>
  5.  
  6. using namespace std;
  7.  
  8. vector<vector<pair<int, int>>> g;
  9. vector<int> parent, edgeUp, depth;
  10.  
  11. void rootDFS(int node, int d = 0, int eUP = -1, int par = -1) {
  12.   parent[node] = par;
  13.   edgeUp[node] = eUP;
  14.   depth[node] = d;
  15.   for (auto [child, e] : g[node]) {
  16.     if (child != par) rootDFS(child, d+1, e, node);
  17.   }
  18. }
  19. void rootTree(int root, int N) {
  20.   parent.resize(N);
  21.   edgeUp.resize(N);
  22.   depth.resize(N);
  23.   rootDFS(root);
  24. }
  25. vector<int> nodesDepth(int d) {
  26.   vector<int> v;
  27.   for (int i = 0; i < (int)depth.size(); ++i) {
  28.     if (depth[i] == d) v.push_back(i);
  29.   }
  30.   return v;
  31. }
  32.  
  33. void find_pair(int N, vector<int> U, vector<int> V, int A, int B) {
  34.   g.resize(N);
  35.   int M = U.size();
  36.   for (int i = 0; i < M; ++i) {
  37.     g[U[i]].emplace_back(V[i], i);
  38.     g[V[i]].emplace_back(U[i], i);
  39.   }
  40.   vector<int> w(M);
  41.   long long toll = ask(w);
  42.  
  43.   rootTree(0, N);
  44.   // Find the depth of the lower of the two nodes
  45.   int l = 0, r = *max_element(depth.begin(), depth.end()); // l < d <= r
  46.   while (l+1 < r) {
  47.     int mid = (l+r)/2;
  48.     // set all nodes lower than depth mid to have B edges leading up
  49.     for (int i = 0; i < N; ++i) {
  50.       if (depth[i] > mid) w[edgeUp[i]] = 1;
  51.     }
  52.     long long res = ask(w);
  53.     for (int i = 0; i < N; ++i) {
  54.       if (depth[i] > mid) w[edgeUp[i]] = 0;
  55.     }
  56.     if (res == toll) r = mid;
  57.     else l = mid;
  58.   }
  59.   int lowerDepth = r;
  60.  
  61.   auto findWithDepth = [&toll, &w](int depth) {
  62.     vector<int> candidates = nodesDepth(depth);
  63.     int l = 0, r = (int)candidates.size(); // in range [l, r)
  64.     while (l+1 < r) {
  65.       int mid = (l+r)/2;
  66.       for (int i = mid; i < r; ++i) w[edgeUp[candidates[i]]] = 1, assert(edgeUp[candidates[i]] >= 0);
  67.       long long res = ask(w);
  68.       for (int i = mid; i < r; ++i) w[edgeUp[candidates[i]]] = 0;
  69.       if (res == toll) r = mid;
  70.       else l = mid;
  71.     }
  72.     return candidates[l];
  73.   };
  74.  
  75.   // Find the lower of the two nodes (or one of them if they have the same depth)
  76.   int S = findWithDepth(lowerDepth);
  77.  
  78.   // Find the other node by rooting the tree at the first node
  79.   rootTree(S, N);
  80.   assert(toll%A==0);
  81.   int T = findWithDepth(toll/A);
  82.   assert(depth[T] == toll/A);
  83.   answer(S, T);
  84. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement