Submission #499687

#	Time	Username	Problem	Language	Result	Execution time	Memory
499687		CatalinT	Highway Tolls (IOI18_highway)	C++17	36 / 100	592 ms	26000 KiB

highway

/*
TODO: Shuffle everything
*/

#include <map>
#include <unordered_map>
#include <unordered_set>
#include <set>
#include <vector>
#include <numeric>
#include <algorithm>
#include <iostream>
#include <string>
#include <cstring>
#include <sstream>
#include <functional>
#include <queue>
#include <deque>
#include <stack>
#include <cassert>
#include <iomanip>
#include <cmath>
#include <bitset>
#include <random>

using namespace std;

#include "highway.h"

using int64 = long long;

// cache

map<pair<int, int>, int> eidx;

int eget(int u, int v) {
    if (u > v) swap(u, v);
    return eidx[{u, v}];
}

void eset(int u, int v, int idx) {
    if (u > v) swap(u, v);
    eidx[{u, v}] = idx;
}

//

vector<vector<int>> adj;

vector<int> tree_idx;
vector<int> comp;
vector<int> dist;
vector<int> par;

void find_pair(int N, std::vector<int> U, std::vector<int> V, int A, int B) {
    int M = U.size();
    int S = -1;
    int T = -1;
    // Step 0. Find min dist, 1 query

    random_device rd;
    mt19937 rnd(rd());
 

    vector<int> w(M);
    int64 D = ask(w) / A;

    // Step 1. Find edge on min - path, 18 queries worst-case (!)

    int l = 0, r = M - 1, sol = -1;

    vector<int> order(M);
    iota(order.begin(), order.end(), 0);

    // shuffle(order.begin(), order.end(), rnd);

    while (l <= r) {
        int m = (l + r) >> 1;

        w.assign(M, 0);
        for (int i = 0; i <= m; i++)
            w[order[i]] = 1;

        int64 cur = ask(w);

        if (cur > D * A) {
            sol = m;
            r = m - 1;
        } else {
            l = m + 1;
        }
    }

    assert(sol != -1);

    sol = order[sol];

    cerr << "Critical edge: " << U[sol] << " -> " << V[sol] << "\n";
    cerr << "D: " << D << "\n";

    adj.assign(N, vector<int>{});
    
    for (int i = 0; i < M; i++) {
        adj[ U[i] ].push_back( V[i] );
        adj[ V[i] ].push_back( U[i] );
        eset(U[i], V[i], i);
    }

    const int x = U[sol];
    const int y = V[sol];

    queue<int> Q;

    Q.push(x);
    Q.push(y);

    dist.assign(N, -1);
    par.assign(N, -1);
    comp.assign(N, -1);

    comp[x] = 0;
    comp[y] = 1;

    dist[x] = 0;
    dist[y] = 0;

    tree_idx.push_back(eget(x, y));

    for (; !Q.empty(); ) {
        auto el = Q.front();
        Q.pop();

        for (auto v : adj[el]) {
            if (dist[v] == -1) {
                dist[v] = dist[el] + 1;
                comp[v] = comp[el];
                par[v] = el;
                tree_idx.push_back(eget(el, v));
                Q.push(v);
            }
        }
    }

    assert(tree_idx.size() == N - 1);

    // TODO: Remove
    // w.assign(M, 1);
    // for (auto e : tree_idx) {
    //     w[e] = 0;
    // }
    // assert(ask(w) == D * A);
    //

    // w.assign(M, 1);
    // for (int i = 0; i < N; i++) {
    //     if (par[i] != -1 && comp[i] == 1) {
    //         w[eget(i, par[i])] = 0;
    //     }
    // }

    // int64 depth = ask(w);

    // bool ok = false;
    // for (int64 i = 0; i < N; i++) {
    //     if (i * A + (D - i) * B == depth) {
    //         depth = i;
    //         ok = true;
    //         break;
    //     }
    // }
    // assert(ok);

    // depth = D - 1 - depth;
    int depth = 0;

    l = 0;
    r = 0;
    sol = 0;

    for (int i = 0; i < N; i++) {
        if (comp[i] == 0 && par[i] != -1) {
            r = max(r, dist[i]);
        }
    }

    if (!r) {
        depth = 0;
    } else {
        sol = r;
        r--;
        while(l <= r) {
            int m = (l + r) >> 1;

            w.assign(M, 1);
            for (auto e : tree_idx)
                w[e] = 0;
            for (int i = 0; i < N; i++) {
                if (comp[i] == 0 && par[i] != -1 && dist[i] > m) {
                    w[eget(i, par[i])] = 1;
                }
            }

            if (ask(w) == D * A) {
                sol = m;
                r = m - 1;
            } else {
                l = m + 1;
            }
        }

        assert(sol != -1);
        depth = sol;
    }

    auto solve_half_tree = [&] (int x, int y, int cc, int depth, int & R) {
        R = -1;

        cerr << "solve_half_tree, start: " << x << " par: " << y << " " << cc << "\n";
        cerr << "depth: " << depth << "\n";

        if (!depth) {
            R = x;
            assert(!dist[R]);
        } else {
            vector<int> cand;
            for (int i = 0; i < N; i++) {
                if (par[i] != -1 && dist[i] == depth && comp[i] == cc) {
                    cand.push_back(i);
                }
            }
            assert(cand.size());

            // shuffle(cand.begin(), cand.end(), rnd);

            function<int(int, int)> rec = [&] (int l, int r) {
                if (l == r) {
                    return cand[l];
                }

                int m = (l + r) >> 1;

                vector<int> w(M, 1);

                for (auto e : tree_idx)
                    w[e] = 0;

                for (int i = l; i <= m; i++) {
                    w[eget(cand[i], par[cand[i]])] = 1;
                }

                bool ok = ask(w) > D * A;

                if (ok) {
                    return rec(l, m);
                } else {
                    return rec(m + 1, r);
                }
            };

            R = rec(0, cand.size() - 1);

            assert(dist[R] == depth);
        };
    };

    solve_half_tree(x, y, 0, depth, S);
    solve_half_tree(y, x, 1, D - 1 - depth, T);

    cerr << S << " " << T << "\n";

    // auto debug = [&](int _S, int _T) {
    //     cerr << dist[_S] << " " << comp[_S] << "\n";
    //     cerr << dist[_T] << " " << comp[_T] << "\n";
    // };

    // cerr << "me:\n";
    // debug(S, T);
    // cerr << "correct:\n";
    // debug(89669, 73110);
    //debug(RS, RT);

    assert(S != T);

    answer(S, T);
}

Compilation message (stderr)

In file included from /usr/include/c++/10/cassert:44,
                 from highway.cpp:20:
highway.cpp: In function 'void find_pair(int, std::vector<int>, std::vector<int>, int, int)':
highway.cpp:144:28: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  144 |     assert(tree_idx.size() == N - 1);
      |            ~~~~~~~~~~~~~~~~^~~~~~~~

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