Submission #995463

#	Time	Username	Problem	Language	Result	Execution time	Memory
995463		shiomusubi496	Highway Tolls (IOI18_highway)	C++17	6 / 100	85 ms	11576 KiB

highway

#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 (stderr)

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);
      |            ~~~~~~~~~^~~~~~~~

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