Submission #766372

#	Time	Username	Problem	Language	Result	Execution time	Memory
766372		boris_mihov	Highway Tolls (IOI18_highway)	C++17	51 / 100	172 ms	262144 KiB

highway

#include "highway.h"
#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <vector>
#include <queue>

typedef long long llong;
const int MAXN = 90000 + 10;
const int INF  = 1e9;

int n, m;
int sz[MAXN];
int par[MAXN];
int parEdge[MAXN];
std::queue <int> q;
std::vector <int> w;
std::vector <std::pair <int,int>> g[MAXN];
std::queue <std::pair <int,int>> qDist;

void initDFS(int node, int p)
{
    sz[node] = 1;
    for (const auto &[u, idx] : g[node])
    {
        if (u == p)
        {
            continue;
        }   

        par[u] = node;
        parEdge[u] = idx;
        initDFS(u, node);
        sz[node] += sz[u];
    }

    if (node == 1)
    {
        sz[node]--;
    }
}

int markEdges(int root, int cnt)
{
    std::fill(w.begin(), w.end(), 0);
    if (cnt == 0)
    {
        return -1;
    }

    while (!q.empty())
    {
        q.pop();
    }

    q.push(root);
    int last = -1;
    if (root != 1)
    {
        w[parEdge[root]] = 1;
        last = parEdge[root];
        cnt--;
    }

    while (cnt)
    {
        int top = q.front();
        q.pop();

        for (const auto &[u, idx] : g[top])
        {
            if (cnt == 0)
            {
                break;
            }

            if (u == par[top])
            {
                continue;
            }

            w[idx] = 1;
            last = idx;
            q.push(u);
            cnt--;
        }
    }

    return last;
}

std::vector <std::pair <int,int>> getKth(int node, int k)
{
    assert(node != 1);
    w[parEdge[node]] = 1;
    std::vector <std::pair <int,int>> candidates;
    qDist.push({node, 1});

    while (!qDist.empty())
    {
        auto [top, dist] = qDist.front();
        qDist.pop();

        if (dist == k)
        {
            candidates.push_back({top, parEdge[top]});
            continue;
        }

        for (const auto &[u, idx] : g[top])
        {
            if (u == par[top])
            {
                continue;
            }

            w[idx] = 1;
            qDist.push({u, dist + 1});
        }
    }

    return candidates;
}

void find_pair(int N, std::vector <int> U, std::vector <int> V, int A, int B) 
{ 
    n = N;
    m = U.size();

    for (int i = 0 ; i < m ; ++i)
    {
        U[i]++;
        V[i]++;
        g[U[i]].push_back({V[i], i});
        g[V[i]].push_back({U[i], i});
    }

    initDFS(1, 0);
    w.resize(m, 0);
    llong dist = ask(w) / A;
    int l = 0, r = n, mid, last = -1;
    while (l < r - 1)
    {
        mid = (l + r) / 2;
        int res = markEdges(1, mid);
        if (ask(w) == dist * A) l = mid;
        else
        {
            r = mid;
            last = res;
        }
    }

    assert(last != -1);
    if (U[last] == par[V[last]])
    {
        std::swap(U[last], V[last]);
    }

    markEdges(U[last], sz[U[last]]);
    llong res = ask(w);
    llong distX = (res - dist * A) / (B - A);
    llong distY = dist - distX;

    std::fill(w.begin(), w.end(), 0);
    std::vector <std::pair <int,int>> candidates = getKth(U[last], distX);

    l = -1;
    r = candidates.size() - 1;

    while (l < r - 1)
    {
        mid = (l + r) / 2;
        for (const auto &[u, idx] : candidates)
        {
            w[idx] = 0; 
        }

        for (int i = 0 ; i <= mid ; ++i)
        {
            w[candidates[i].second] = 1;
        }
    
        res = ask(w);
        if (res == distX * B + distY * A) r = mid;
        else l = mid;
    }

    int s, t;
    s = candidates[r].first;
    if (distY == 0)
    {
        answer(s - 1, V[last] - 1);
        return;
    }

    candidates.clear();
    bool isGood = false;
    std::fill(w.begin(), w.end(), 0);
    for (const auto &[u, idx] : g[V[last]])
    {
        if (u == par[V[last]])
        {
            continue;
        }

        if (u == U[last])
        {
            isGood = true;
            continue;
        }

        if (isGood)
        {
            std::vector <std::pair <int,int>> curr = getKth(u, distY);
            for (const auto &i : curr)
            {
                candidates.push_back(i);
            }
        }
    }

    l = -1, r = candidates.size() - 1;
    while (l < r - 1)
    {
        mid = (l + r) / 2;
        for (const auto &[u, idx] : candidates)
        {
            w[idx] = 0; 
        }

        for (int i = 0 ; i <= mid ; ++i)
        {
            w[candidates[i].second] = 1;
        }
    
        res = ask(w);
        if (res == distX * A + distY * B) r = mid;
        else l = mid;
    }

    t = candidates[r].first;
    answer(s - 1, t - 1);
}   

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