#include "highway.h"
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define rep2(i, a, b) for (int i = (int)(a); i < (int)(b); ++i)
#define rrep(i, n) for (int i = (int)(n) - 1; i >= 0; --i)
#define rrep2(i, a, b) for (int i = (int)(b) - 1; i >= (int)(a); --i)
#define all(v) begin(v), end(v)
using namespace std;
using ll = long long;
template<class T, class U> bool chmin(T& a, const U& b) { return a > b ? a = b, true : false; }
template<class T, class U> bool chmax(T& a, const U& b) { return a < b ? a = b, true : false; }
mt19937 mt(random_device{}());
struct edge {
int to, idx;
edge() = default;
edge(int t, int i) : to(t), idx(i) {}
};
void find_pair(int N, std::vector<int> U, std::vector<int> V, int A, int B) {
assert(U.size() == N - 1);
vector<vector<edge>> G(N);
rep (i, N - 1) {
G[U[i]].emplace_back(V[i], i);
G[V[i]].emplace_back(U[i], i);
}
vector<bool> used(N);
vector<int> sz(N);
auto find_centroid = [&](int v) {
fill(all(sz), 0);
auto dfs = [&](auto&& dfs, int v, int p) -> void {
sz[v] = 1;
for (auto e : G[v]) if (e.to != p && !used[e.to]) {
dfs(dfs, e.to, v);
sz[v] += sz[e.to];
}
};
dfs(dfs, v, -1);
int n = sz[v];
auto dfs2 = [&](auto&& dfs2, int v, int p) -> int {
for (auto e : G[v]) if (e.to != p && !used[e.to] && sz[e.to] * 2 > n) {
return dfs2(dfs2, e.to, v);
}
return v;
};
return dfs2(dfs2, v, -1);
};
ll dist = ask(vector<int>(N - 1, 0)) / A;
auto query = [&](vector<int> f) {
return (ask(f) - dist * A) / (B - A);
};
if (dist == 1) {
int ng = 0, ok = N - 1;
while (ok - ng > 1) {
int mid = (ok + ng) / 2;
vector<int> f(N - 1);
rep (i, N - 1) f[i] = i < mid;
if (query(f)) ok = mid;
else ng = mid;
}
answer(U[ok - 1], V[ok - 1]);
return;
}
int cur_v = 0;
while (1) {
auto v = find_centroid(cur_v);
find_centroid(v);
used[v] = true;
vector<pair<int, int>> sub;
for (auto e : G[v]) if (!used[e.to]) {
sub.emplace_back(sz[e.to], e.to);
}
shuffle(all(sub), mt);
int sm = 0;
vector<bool> memo(N);
for (auto [s, u] : sub) {
if (sm + s <= sz[v] * 2 / 3) {
sm += s;
memo[u] = true;
}
}
vector<int> f(N - 1);
for (auto e : G[v]) if (memo[e.to]) {
f[e.idx] = 1;
auto dfs = [&](auto&& dfs, int v, int p) -> void {
for (auto e : G[v]) if (e.to != p && !used[e.to]) {
f[e.idx] = 1;
dfs(dfs, e.to, v);
}
};
dfs(dfs, e.to, v);
}
int d = query(f);
if (d == dist) {
for (auto e : G[v]) {
if (memo[e.to]) {
cur_v = e.to;
}
else {
used[e.to] = true;
}
}
}
else if (d == 0) {
for (auto e : G[v]) {
if (!memo[e.to]) {
cur_v = e.to;
}
else {
used[e.to] = true;
}
}
}
else {
vector<int> dep(N);
auto dfs = [&](auto&& dfs, int v, int p) -> void {
for (auto e : G[v]) if (e.to != p && !used[e.to]) {
dep[e.to] = dep[v] + 1;
dfs(dfs, e.to, v);
}
};
dfs(dfs, v, -1);
vector<int> X, Y;
for (auto e : G[v]) {
bool f = memo[e.to];
int x = f ? d : dist - d;
auto dfs = [&](auto&& dfs, int v, int p) -> void {
for (auto e : G[v]) if (e.to != p && !used[e.to]) {
if (dep[e.to] == x) {
if (f) X.push_back(e.idx);
else Y.push_back(e.idx);
}
dfs(dfs, e.to, v);
}
};
if (x == 1) {
if (f) X.push_back(e.idx);
else Y.push_back(e.idx);
}
dep[e.to] = 1;
dfs(dfs, e.to, v);
}
int ans1, ans2;
{
int ng = 0, ok = X.size();
while (ok - ng > 1) {
int mid = (ok + ng) / 2;
vector<int> f(N - 1);
rep (i, mid) f[X[i]] = 1;
if (query(f)) ok = mid;
else ng = mid;
}
ans1 = dep[U[X[ok - 1]]] == d ? U[X[ok - 1]] : V[X[ok - 1]];
}
{
int ng = 0, ok = Y.size();
while (ok - ng > 1) {
int mid = (ok + ng) / 2;
vector<int> f(N - 1);
rep (i, mid) f[Y[i]] = 1;
if (query(f)) ok = mid;
else ng = mid;
}
ans2 = dep[U[Y[ok - 1]]] == dist - d ? U[Y[ok - 1]] : V[Y[ok - 1]];
}
answer(ans1, ans2);
return;
}
}
}
Compilation message
In file included from /usr/include/c++/10/cassert:44,
from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:33,
from highway.cpp:2:
highway.cpp: In function 'void find_pair(int, std::vector<int>, std::vector<int>, int, int)':
highway.cpp:26:21: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
26 | assert(U.size() == N - 1);
| ~~~~~~~~~^~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
1 ms |
344 KB |
Execution killed with signal 11 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
344 KB |
Output is correct |
| 2 |
Correct |
10 ms |
1456 KB |
Output is correct |
| 3 |
Correct |
73 ms |
9120 KB |
Output is correct |
| 4 |
Correct |
84 ms |
9176 KB |
Output is correct |
| 5 |
Incorrect |
85 ms |
9164 KB |
Output is incorrect: {s, t} is wrong. |
| 6 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
6 ms |
1624 KB |
Output is correct |
| 2 |
Correct |
16 ms |
2904 KB |
Output is correct |
| 3 |
Correct |
16 ms |
4264 KB |
Output is correct |
| 4 |
Correct |
68 ms |
11520 KB |
Output is correct |
| 5 |
Correct |
63 ms |
11556 KB |
Output is correct |
| 6 |
Correct |
72 ms |
11524 KB |
Output is correct |
| 7 |
Correct |
66 ms |
11576 KB |
Output is correct |
| 8 |
Correct |
72 ms |
11504 KB |
Output is correct |
| 9 |
Correct |
61 ms |
11520 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
1 ms |
600 KB |
Execution killed with signal 11 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
6 ms |
856 KB |
Execution killed with signal 6 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
7 ms |
856 KB |
Execution killed with signal 6 |
| 2 |
Halted |
0 ms |
0 KB |
- |
