Submission #691868

#	Time	Username	Problem	Language	Result	Execution time	Memory
691868		piOOE	Prize (CEOI22_prize)	C++17	100 / 100	2361 ms	343084 KiB

Main

#include <bits/stdc++.h>

using namespace std;

#define int int64_t

void massert(bool v) {
    if (!v) {
        auto start = clock();
        int x = 1;
        while ((clock() - start) < 5 * CLOCKS_PER_SEC) {
            x <<= 1;
            x += 1;
        }
    }
}

struct HLD {
    vector<int> par, siz, head, tin, tout, ord, depth;

    int dfs1(int i, vector<vector<int>> &g) {
        for (int &t: g[i]) {
            if (t != par[i]) {
                depth[t] = depth[i] + 1;
                par[t] = i;
                siz[i] += dfs1(t, g);
                if (siz[t] > siz[g[i][0]] || g[i][0] == par[i]) swap(t, g[i][0]);
            }
        }
        return siz[i];
    }

    void dfs2(int i, int &T, const vector<vector<int>> &g) {
        tin[i] = T++;
        for (int t: g[i]) {
            if (t != par[i]) {
                head[t] = (T == tin[i] + 1 ? head[i] : t);
                dfs2(t, T, g);
            }
        }
        tout[i] = T;
    }

    HLD() = default;

    HLD(vector<vector<int>> g, int r = 0)
            : par(g.size(), -1), siz(g.size(), 1), head(g.size(), r), tin(g.size()), ord(g.size()), tout(g.size()),
              depth(g.size()) {
        dfs1(r, g);
        int x = 0;
        dfs2(r, x, g);
        for (int i = 0; i < g.size(); ++i) ord[tin[i]] = i;
    }

    vector<pair<int, int>> path(int a, int b) {
        vector<pair<int, int>> res;
        for (;; b = par[head[b]]) {
            if (tin[b] < tin[a]) swap(a, b);
            if (tin[head[b]] <= tin[a]) {
                res.emplace_back(tin[a], tin[b] + 1);
                return res;
            }
            res.emplace_back(tin[head[b]], tin[b] + 1);
        }
    }

    pair<int, int> subtree(int i) {
        return {tin[i], tin[i] + siz[i]};
    }

    int dist(int a, int b) {
        return depth[a] + depth[b] - 2 * depth[lca(a, b)];
    }

    int lca(int a, int b) {
        for (;; b = par[head[b]]) {
            if (tin[b] < tin[a]) swap(a, b);
            if (tin[head[b]] <= tin[a]) return a;
        }
    }

    bool isp(int a, int b) {
        return tin[a] <= tin[b] && tout[a] >= tout[b];
    }
} t1, t2;

vector<int> normalize(int n, int root, vector<vector<pair<int, int>>> g) {
    vector<int> dist(n, -1);
    dist[root] = 0;
    vector<bool> used(n);

    function<void(int)> dfs = [&](int v) {
        used[v] = true;
        for (auto [to, w]: g.at(v)) {
            if (!used[to]) {
                dist[to] = dist[v] + w;
                dfs(to);
            }
        }
    };

    dfs(root);

    return dist;
}

int32_t main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, k, q, t;
    cin >> n >> k >> q >> t;

    vector<int> p1(n), p2(n);

    int root1 = -1, root2 = -1;
    for (int i = 0; i < n; ++i) {
        cin >> p1[i];
        if (p1[i] == -1) {
            root1 = i;
        } else {
            p1[i] -= 1;
        }
    }

    for (int i = 0; i < n; ++i) {
        cin >> p2[i];
        if (p2[i] == -1) {
            root2 = i;
        } else {
            p2[i] -= 1;
        }
    }

    massert(max(root1, root2) < n && min(root1, root2) >= 0);

    vector<vector<int>> adj(n);

    for (int i = 0; i < n; ++i) {
        if (p1[i] != -1) {
            adj[p1[i]].push_back(i);
        }
    }

    t1 = HLD(adj, root1);

    for (int i = 0; i < n; ++i) {
        adj[i].clear();
    }

    for (int i = 0; i < n; ++i) {
        if (p2[i] != -1) {
            adj[p2[i]].push_back(i);
        }
    }

    t2 = HLD(adj, root2);

    vector<pair<int, int>> questions;

    vector<int> ord(n);
    iota(ord.begin(), ord.end(), 0);
    sort(ord.begin(), ord.end(), [&](int i, int j) {
        return t1.tin[i] < t1.tin[j];
    });

    ord.resize(k);
    
    for (int x : ord) {
        cout << x + 1 << " \n"[x == ord.back()];
    }
 
    sort(ord.begin(), ord.end(), [&](int i, int j) {
        return t2.tin[i] < t2.tin[j];
    });

    for (int i = 1; i < ord.size(); ++i) {
        questions.emplace_back(ord[i - 1], ord[i]);
    }

    assert(questions.size() <= q);
    for (auto [x, y]: questions) {
        cout << "? " << x + 1 << " " << y + 1 << '\n';
    }
    cout << "!" << endl;

    vector<vector<pair<int, int>>> g1(n), g2(n);

    int h1 = ord[0], h2 = ord[0];

    vector<int> aa, bb;
    auto addEdge = [&](int tt, int x, int y, int w) {
        if (x == y) {
            return;
        }
        if (tt == 0) aa.push_back(x), aa.push_back(y);
        else bb.push_back(x), bb.push_back(y);
        massert(x < n && y < n && x >= 0 && y >= 0);
        (tt == 0 ? g1[x] : g2[x]).emplace_back(y, w);
        (tt == 0 ? g1[y] : g2[y]).emplace_back(x, -w);
    };

    for (auto [x, y]: questions) {
        int a, b, c, d;
        cin >> a >> b >> c >> d;

        int c1 = t1.lca(x, y);
        int c2 = t2.lca(x, y);

        if (t1.isp(c1, h1)) {
            h1 = c1;
        }
        if (t2.isp(c2, h2)) {
            h2 = c2;
        }

        addEdge(0, c1, x, a);
        addEdge(0, c1, y, b);
        addEdge(1, c2, x, c);
        addEdge(1, c2, y, d);
    }

    auto first = normalize(n, h1, g1);
    auto second = normalize(n, h2, g2);

    vector<pair<int, int>> queries(t);
    for (auto &[a, b]: queries) {
        cin >> a >> b;
        a -= 1, b -= 1;
    }

    for (int i: aa) {
        assert(first[i] != -1);
    }
    for (int i: bb) {
        assert(second[i] != -1);
    }

    for (auto [a, b]: queries) {
        int dist1 = first[a] - 2 * first[t1.lca(a, b)] + first[b];
        int dist2 = second[a] - 2 * second[t2.lca(a, b)] + second[b];

        cout << dist1 << " " << dist2 << '\n';
    }

    return 0;
}

Compilation message (stderr)

Main.cpp: In constructor 'HLD::HLD(std::vector<std::vector<long int> >, int64_t)':
Main.cpp:19:44: warning: 'HLD::ord' will be initialized after [-Wreorder]
   19 |     vector<int> par, siz, head, tin, tout, ord, depth;
      |                                            ^~~
Main.cpp:19:38: warning:   'std::vector<long int> HLD::tout' [-Wreorder]
   19 |     vector<int> par, siz, head, tin, tout, ord, depth;
      |                                      ^~~~
Main.cpp:46:5: warning:   when initialized here [-Wreorder]
   46 |     HLD(vector<vector<int>> g, int r = 0)
      |     ^~~
Main.cpp:52:27: warning: comparison of integer expressions of different signedness: 'int64_t' {aka 'long int'} and 'std::vector<std::vector<long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   52 |         for (int i = 0; i < g.size(); ++i) ord[tin[i]] = i;
      |                         ~~^~~~~~~~~~
Main.cpp: In function 'int32_t main()':
Main.cpp:177:23: warning: comparison of integer expressions of different signedness: 'int64_t' {aka 'long int'} and 'std::vector<long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  177 |     for (int i = 1; i < ord.size(); ++i) {
      |                     ~~^~~~~~~~~~~~
In file included from /usr/include/c++/10/cassert:44,
                 from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:33,
                 from Main.cpp:1:
Main.cpp:181:29: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<long int, long int> >::size_type' {aka 'long unsigned int'} and 'int64_t' {aka 'long int'} [-Wsign-compare]
  181 |     assert(questions.size() <= q);
      |            ~~~~~~~~~~~~~~~~~^~~~

• Subtask 1: 10 / 10
• Subtask 2: 25 / 25
• Subtask 3: 19 / 19
• Subtask 4: 22 / 22
• Subtask 5: 24 / 24