Submission #951682

#	Time	Username	Problem	Language	Result	Execution time	Memory
951682		kilkuwu	Stations (IOI20_stations)	C++17	100 / 100	600 ms	1484 KiB

stations

#ifndef LOCAL
#include "stations.h"
#else
#include <vector>

std::vector<int> label(int n, int k, std::vector<int> u, std::vector<int> v);
int find_next_station(int s, int t, std::vector<int> c);

#include <cstdio>
#include <cassert>
#include <map>
#include <vector>
#include <algorithm>

static int __max_label = 0;
static int __r, __n, __k, __q;
static std::vector<int> __u, __v, __labels, __answers;
static std::map<int, int> __reverse_labels;
static std::vector<std::vector<int>> __adjlist;
static int __s, __t, __w;
static std::vector<int> __c;

int main() {
	assert(scanf("%d", &__r) == 1);
	for (int tc = 0; tc < __r; tc++) {
		assert(scanf("%d%d", &__n, &__k) == 2);
		__u.resize(__n - 1);
		__v.resize(__n - 1);
		__adjlist.clear();
		__adjlist.resize(__n);
		for (int i = 0; i < __n - 1; i++) {
			assert(scanf("%d%d", &__u[i], &__v[i]) == 2);
			__adjlist[__u[i]].push_back(__v[i]);
			__adjlist[__v[i]].push_back(__u[i]);
		}
		__labels = label(__n, __k, __u, __v);
		if ((int)__labels.size() != __n) {
			printf("Number of labels not equal to %d\n", __n);
			exit(0);
		}
		__reverse_labels.clear();
		for (int i = 0; i < __n; i++) {
			if (__labels[i] < 0 || __labels[i] > __k) {
				printf("Label not in range 0 to %d\n", __k);
				exit(0);
			}
			if (__reverse_labels.find(__labels[i]) != __reverse_labels.end()) {
				printf("Labels not unique\n");
				exit(0);
			}
			__reverse_labels[__labels[i]] = i;
			if (__labels[i] > __max_label) {
				__max_label = __labels[i];
			}
		}
		assert(scanf("%d", &__q) == 1);
		for (int i = 0; i < __q; i++) {
			assert(scanf("%d%d%d", &__s, &__t, &__w) == 3);
			__c.clear();
			for (int v : __adjlist[__s]) {
				__c.push_back(__labels[v]);
			}
			std::sort(__c.begin(), __c.end());
			int answer = find_next_station(__labels[__s], __labels[__t], __c);
			if (!std::binary_search(__c.begin(), __c.end(), answer)) {
				printf("Label %d returned by find_next_station not found in c\n", answer);
				exit(0);
			}
			__answers.push_back(__reverse_labels[answer]);
		}
	}
	printf("%d\n", __max_label);
	for (int index : __answers) {
		printf("%d\n", index);
	}
	exit(0);
}
#endif
#include <bits/stdc++.h>

std::vector<int> label(int n, int k, std::vector<int> u, std::vector<int> v) {
  std::vector<std::vector<int>> adj(n);
  std::vector<int> st(n), en(n), dep(n);
  for (int i = 0; i + 1 < n; i++) {
    adj[u[i]].push_back(v[i]); 
    adj[v[i]].push_back(u[i]);
  }

  int ti = 0;
  std::function<void(int, int)> dfs = [&](int x, int p) -> void {
    st[x] = ti++;
    for (int y : adj[x]) {
      if (y == p) continue;
      dep[y] = dep[x] + 1;
      dfs(y, x);
    }
    en[x] = ti++;
  };

  dfs(0, 0);

  std::vector<int> res(n);
  for (int i = 0; i < n; i++) {
    if (dep[i] & 1) {
      res[i] = en[i];
    } else {
      res[i] = st[i];
    }
  }

  std::vector<int> vals = res;
  std::sort(vals.begin(), vals.end());
  for (int i = 0; i < n; i++) {
    res[i] = std::lower_bound(vals.begin(), vals.end(), res[i]) - vals.begin();
  }

  return res;
}

int find_next_station(int s, int t, std::vector<int> c) {
  int sz = c.size();
  if (s == 0) {
    // we know this guy the root 
    int st = 1;
    for (int i = 0; i < sz; i++) {
      if (st <= t && t <= c[i]) {
        return c[i];
      }
      st = c[i] + 1;
    }
    return c.back();
  }
  
  if (s < c[0]) {
    // this means this guy is st 

    // biggest is the out

    int st = s + 1;
    for (int i = 0; i + 1 < sz; i++) {
      if (st <= t && t <= c[i]) {
        return c[i];
      }
      st = c[i] + 1;
    }
    
    return c.back();
  } else {
    int en = s - 1;
    for (int i = sz - 1; i > 0; i--) {
      if (c[i] <= t && t <= en) {
        return c[i];
      }
      en = c[i] - 1;
    }
    return c[0];
  }
}

• Subtask 1: 5 / 5
• Subtask 2: 8 / 8
• Subtask 3: 16 / 16
• Subtask 4: 10 / 10
• Subtask 5: 61 / 61