Submission #390834

#	Time	Username	Problem	Language	Result	Execution time	Memory
390834		JimmyZJX	Highway Tolls (IOI18_highway)	C++14	51 / 100	346 ms	14600 KiB

highway

#include <iostream>
#include <fstream>
#include <algorithm>
#include <cstring>
#include <climits>
#include <cassert>
#include <tuple>
#include <queue>
#include <stack>
#include <vector>
#include <map>
#include <set>
#include <string>
#include <unordered_set>
#include <unordered_map>
#include <random>

#include "highway.h"

using namespace std;

typedef long long LL;
typedef vector<int> Vi;
typedef vector<bool> Vb;
typedef vector<vector<int>> Vii;

#define forR(i, n) for (int i = 0; i < (n); i++)

#define MAXMN 130003

vector<vector<pair<int, int>>> nbs(MAXMN); // {nb, edgeNo}
vector<int> parent(MAXMN);
vector<int> parentEdge(MAXMN);
Vi edgeSorted;
Vi edgeRank;
int maxDist;

int parentOf(int n1, int n2) {
	if (parent[n1] == n2) return n2;
	return n1;
}

int childOf(int n1, int n2) {
	if (parent[n1] == n2) return n1;
	return n2;
}

int largestEdgeOfRank(int r) {
	forR(i, edgeSorted.size()) {
		if (edgeRank[edgeSorted[i]] > r)
			return i > 0 ? i - 1 : i;
	}
	return edgeSorted.size() - 1;
}

void build_tree(int root) {
	Vb vis(MAXMN); Vb edgeVis(MAXMN);
	edgeSorted.clear(); edgeRank = Vi(MAXMN); maxDist = 0;
	queue<pair<int, int>> q; q.push({ root, 0 }); // nodeNo, level
	vis[root] = true;
	while (!q.empty()) {
		auto p = q.front(); q.pop();
		int x = p.first, level = p.second;
		maxDist = max(maxDist, level);
		for (auto nbP : nbs[x]) {
			int nb = nbP.first, nbEdge = nbP.second;
			if (!vis[nb]) {
				vis[nb] = true;
				q.push({ nb, level + 1 });
				parent[nb] = x; parentEdge[nb] = nbEdge;
			}
			if (!edgeVis[nbEdge]) {
				edgeVis[nbEdge] = true;
				edgeSorted.push_back(nbEdge);
				edgeRank[nbEdge] = level + 1;
			}
		}
	}
}

void find_pair(int N, Vi U, Vi V, int A, int B) {
	int m = U.size();

	forR(i, m) {
		nbs[U[i]].push_back({ V[i], i });
		nbs[V[i]].push_back({ U[i], i });
	}

	auto g = mt19937();
	for (auto& v : nbs) {
		shuffle(v.begin(), v.end(), g);
	}

	build_tree(0);
	int far_node = childOf(U[edgeSorted.back()], V[edgeSorted.back()]);
	build_tree(far_node);
	int far_node2 = childOf(U[edgeSorted.back()], V[edgeSorted.back()]);
	Vi DiameterEdges;
	for (int n = far_node2; n != far_node; n = parent[n]) {
		DiameterEdges.push_back(parentEdge[n]);
	}
	int midEdge = DiameterEdges[(DiameterEdges.size() - 1) / 2];
	int midNode = childOf(U[midEdge], V[midEdge]);

	//build_tree(midNode);
	build_tree((rand() * 2000) % N);
	//build_tree(far_node2);

	LL minCost = ask(Vi(m));
	LL dist = minCost / A;
	// search for range of edges
	int l = 0, r = m - 1;
	while (l < r) {
		int mid = (l + r) / 2;
		Vi ws(m);
		for (int i = 0; i <= mid; i++) ws[edgeSorted[i]] = 1;
		if (ask(ws) == minCost) {
			l = mid + 1;
		} else {
			r = mid;
		}
	}
	int leftEdge = edgeSorted[l];
	int rootNode = parentOf(U[leftEdge], V[leftEdge]);
	int root2 = childOf(U[leftEdge], V[leftEdge]);

	build_tree(rootNode);

	Vi edgeSorted0 = edgeSorted;
	Vi edgeRank0 = edgeRank;

	build_tree(root2);

	Vi edgesFrom0, edgesFrom1;
	for (int e : edgeSorted0)
		if (edgeRank0[e] < edgeRank[e])
			edgesFrom0.push_back(e);
	for (int e : edgeSorted)
		if (edgeRank[e] < edgeRank0[e])
			edgesFrom1.push_back(e);

	l = -1, r = edgesFrom0.size() - 1; // largestEdgeOfRank(dist);
	int edgeCount = edgesFrom0.size();
	while (l < r) {
		int mid = l + (r - l) / 2;
		Vi ws(m);
		for (int i = mid + 1; i <= edgeCount - 1; i++) ws[edgesFrom0[i]] = 1;
		if (ask(ws) == minCost) {
			r = mid;
		} else {
			l = mid + 1;
		}
	}
	int rightChild;
	if (l < 0) {
		rightChild = rootNode;
	} else {
		int rightEdge = edgesFrom0[l];
		rightChild = childOf(U[rightEdge], V[rightEdge]);
	}

	LL distR = 0; Vi REdges;
	for (int n = rightChild; n != rootNode; n = parent[n]) {
		REdges.push_back(parentEdge[n]);
		distR++;
	}
	LL distL = dist - distR - 1;
	if (distL == 0) {
		answer(root2, rightChild);
		return;
	}


	//l = largestEdgeOfRank(distL - 1), r = min(rightRank, largestEdgeOfRank(distL));
	l = 0, r = edgesFrom1.size() - 1;
	edgeCount = edgesFrom1.size();
	while (l < r) {
		int mid = (l + r) / 2;
		Vi ws(m);
		for (int i = mid + 1; i <= edgeCount - 1; i++) ws[edgesFrom1[i]] = 1;
		for (int e : REdges) ws[e] = 0;
		if (ask(ws) == minCost) {
			r = mid;
		} else {
			l = mid + 1;
		}
	}
	int edgeLeft = edgesFrom1[l];
	int leftChild = childOf(U[edgeLeft], V[edgeLeft]);

	answer(leftChild, rightChild);
}

Compilation message (stderr)

highway.cpp: In function 'int largestEdgeOfRank(int)':
highway.cpp:27:38: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   27 | #define forR(i, n) for (int i = 0; i < (n); i++)
      |                                      ^
highway.cpp:49:2: note: in expansion of macro 'forR'
   49 |  forR(i, edgeSorted.size()) {
      |  ^~~~
highway.cpp: In function 'void find_pair(int, Vi, Vi, int, int)':
highway.cpp:103:6: warning: unused variable 'midNode' [-Wunused-variable]
  103 |  int midNode = childOf(U[midEdge], V[midEdge]);
      |      ^~~~~~~

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