Submission #218836

#	Time	Username	Problem	Language	Result	Execution time	Memory
218836		MarcoMeijer	Highway Tolls (IOI18_highway)	C++14	0 / 100	319 ms	262148 KiB

highway

#include "highway.h"
#include <bits/stdc++.h>
using namespace std;

//macros
typedef long long ll;
typedef pair<int, int> ii;
typedef tuple<int, int, int> iii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<iii> viii;
typedef vector<ll> vll;
#define REP(a,b,c) for(int a=int(b); a<int(c); a++)
#define RE(a,c) REP(a,0,c)
#define RE1(a,c) REP(a,1,c+1)
#define REI(a,b,c) REP(a,b,c+1)
#define REV(a,b,c) for(int a=int(c-1); a>=int(b); a--)
#define pb push_back
#define fi first
#define se second
#define sz size()

struct DSU {
  vi p;
  void create(int size) {
    p.resize(size);
    RE(i,size) p[i] = i;
  }
  int get(int i) {return i==p[i]?i:p[i]=get(p[i]);}
  bool same(int i, int j) {return get(i)==get(j);}
  void unite(int i, int j) {
    if(!same(i,j)) {
      i=get(i); j=get(j);
      p[i] = j;
    }
  }
};

const int MX = 1e5;

int n, m, a, b;
vi u, v, w;
set<ii> adj[MX];
int SZ[MX], P[MX];
vi dist;
ll DISTANCE;

void dfsP(int u=0, int p=-1) {
  P[u] = p;
  for(ii v:adj[u]) if(v.fi != p) dfsP(v.fi,u);
}
void dfsSize(int u=0, int p=-1) {
  SZ[u] = 1;
  for(ii v:adj[u]) if(v.fi != p) dfsSize(v.fi,u), SZ[u]+=SZ[v.fi];
}
int findCentroid(int mxSize, int u, int p=-1) {
  for(ii v:adj[u]) if(v.fi != p) if(SZ[v.fi] > mxSize) return findCentroid(mxSize, v.fi, u);
  return u;
}
void removeEdge(int i) {
  adj[u[i]].erase({v[i],i});
  adj[v[i]].erase({u[i],i});
}
void addEdge(int i) {
  adj[u[i]].insert({v[i],i});
  adj[v[i]].insert({u[i],i});
}
void setWSubTree(int u, int p, int val) {
  for(ii v:adj[u]) if(v.fi != p) setWSubTree(v.fi,u,val), w[v.se] = val;
}
int findS(int u) {
  dfsSize(u, -1);
  if(SZ[u] == 2) {
    if(P[u] == adj[u].begin()->fi) u = P[u];
    w[adj[u].begin()->se] = 1;
    int res = ask(w);
    w[adj[u].begin()->se] = 0;
    if(res == DISTANCE) return u;
    return adj[u].begin()->fi;
  }
  if(SZ[u] == 1) return u;
  int cent = findCentroid(SZ[u]/2, u);
  dfsSize(cent, -1);

  int children;
  vii chl;
  for(ii v:adj[cent]) chl.pb(v);
  DSU dsu;
  dsu.create(chl.sz);
  priority_queue<ii> pq;
  RE(i,chl.sz) pq.push({SZ[chl[i].fi], i});
  while(pq.size() > 2) {
    ii p1 = pq.top(); pq.pop();
    ii p2 = pq.top(); pq.pop();
    dsu.unite(p1.se, p2.se);
    pq.push({p1.fi+p2.fi,dsu.get(p1.se)});
  }
  int subR[2];
  RE(i,2) subR[i] = dsu.get(pq.top().se), pq.pop();

  RE(i,chl.sz) if(dsu.get(i) == subR[0]) {
    w[chl[i].se] = 1;
    setWSubTree(chl[i].fi,cent,1);
  }
  int res = ask(w);
  RE(i,chl.sz) if(dsu.get(i) == subR[0]) {
    w[chl[i].se] = 0;
    setWSubTree(chl[i].fi,cent,0);
  }
  if(res == DISTANCE) {
    // s is in the first sub tree
    RE(i,chl.sz) if(dsu.get(i) == subR[0]) removeEdge(chl[i].se);
    RE(i,chl.sz) if(dsu.get(i) == subR[1]) return findS(chl[i].fi);
  } else {
    // s is in the second sub tree
    RE(i,chl.sz) if(dsu.get(i) == subR[1]) removeEdge(chl[i].se);
    RE(i,chl.sz) if(dsu.get(i) == subR[0]) return findS(chl[i].fi);
  }
}
ii findOnPath(int u) {
  dfsSize(u, -1);
  int cent = findCentroid(SZ[u]/2, u);
  dfsSize(cent, -1);

  int children;
  vii chl;
  for(ii v:adj[cent]) chl.pb(v);
  DSU dsu;
  dsu.create(chl.sz);
  priority_queue<ii> pq;
  RE(i,chl.sz) pq.push({SZ[chl[i].fi], i});
  while(pq.size() > 2) {
    ii p1 = pq.top(); pq.pop();
    ii p2 = pq.top(); pq.pop();
    dsu.unite(p1.se, p2.se);
    pq.push({p1.fi+p2.fi,dsu.get(p1.se)});
  }
  int subR[2];
  RE(i,2) subR[i] = dsu.get(pq.top().se), pq.pop();

  // try both sub trees
  RE(j,2) {
    RE(i,chl.sz) if(dsu.get(i) == subR[j]) {
      w[chl[i].se] = 1;
      setWSubTree(chl[i].fi,cent,1);
    }
    if(ask(w) == DISTANCE) {
      // s and t are in this subTree
      RE(i,chl.sz) if(dsu.get(i) == subR[j]) removeEdge(chl[i].se);
      RE(i,chl.sz) if(dsu.get(i) == subR[!j]) return findOnPath(chl[i].fi);
    }
    RE(i,chl.sz) if(dsu.get(i) == subR[j]) {
      w[chl[i].se] = 0;
      setWSubTree(chl[i].fi,cent,0);
    }
  }

  // s and t are in both sub trees, switch to find s mode
  int ans[2];
  dfsP(cent);
  RE(j,2) {
    RE(i,chl.sz) if(dsu.get(i) == subR[j]) removeEdge(chl[i].se);
    ans[j] = findS(cent);
    RE(i,chl.sz) if(dsu.get(i) == subR[j]) addEdge(chl[i].se);
  }
  return {ans[0],ans[1]};
}

void find_pair(int N, vi U, vi V, int A, int B) {
  n=N; u=U; v=V; a=A; b=B;
  m = U.size();

  RE(i,m) adj[u[i]].insert({v[i],i});
  RE(i,m) adj[v[i]].insert({u[i],i});

  w.assign(m,0);
  DISTANCE = ask(w);

  ii ans = findOnPath(0);

  answer(ans.fi, ans.se);
}

Compilation message (stderr)

highway.cpp: In function 'int findS(int)':
highway.cpp:85:7: warning: unused variable 'children' [-Wunused-variable]
   int children;
       ^~~~~~~~
highway.cpp: In function 'ii findOnPath(int)':
highway.cpp:125:7: warning: unused variable 'children' [-Wunused-variable]
   int children;
       ^~~~~~~~
highway.cpp: In function 'int findS(int)':
highway.cpp:119:1: warning: control reaches end of non-void function [-Wreturn-type]
 }
 ^

• Subtask 1: 0 / 5
• Subtask 2: 0 / 7
• Subtask 3: 0 / 6
• Subtask 4: 0 / 33
• Subtask 5: 0 / 18
• Subtask 6: 0 / 31