Submission #1014517

#	Time	Username	Problem	Language	Result	Execution time	Memory
1014517		shiomusubi496	Prize (CEOI22_prize)	C++17	100 / 100	2199 ms	411964 KiB

Main

#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 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; }

// class LowestCommonAncestor {
//     vector<vector<int>> G;
//     vector<int> dep;
//     vector<vector<int>> par;

//     void dfs(int v, int p) {
//         par[0][v] = p;
//         for (int u : G[v]) {
//             if (u == p) continue;
//             dep[u] = dep[v] + 1;
//             dfs(u, v);
//         }
//     }

// public:
//     void init(vector<vector<int>> G_, int r) {
//         G = G_;
//         int N = G.size();
//         dep.assign(N, 0);
//         par.assign(22, vector<int>(N, -1));
//         dfs(r, -1);
//         rep (i, 21) rep (v, N) {
//             par[i + 1][v] = par[i][v] == -1 ? -1 : par[i][par[i][v]];
//         }
//     }
//     int lca(int a, int b) const {
//         if (dep[a] > dep[b]) swap(a, b);
//         rrep (i, 22) {
//             if ((dep[b] - dep[a]) >> i & 1) b = par[i][b];
//         }
//         if (a == b) return a;
//         rrep (i, 22) {
//             if (par[i][a] != par[i][b]) {
//                 a = par[i][a];
//                 b = par[i][b];
//             }
//         }
//         return par[0][a];
//     }
//     int depth(int v) const { return dep[v]; }
// };


class LowestCommonAncestor {
    class SegmentTree {
        using T = pair<int, int>;
        int n;
        vector<T> dat;

    public:
        void init(vector<T> v) {
            n = 1;
            while (n < (int)v.size()) n <<= 1;
            dat.assign(2 * n, {1e9, -1});
            rep (i, v.size()) dat[i + n] = v[i];
            for (int i = n - 1; i > 0; --i) dat[i] = min(dat[i << 1], dat[i << 1 | 1]);
        }
        T prod(int l, int r) const {
            T res = {1e9, -1};
            l += n; r += n;
            while (l < r) {
                if (l & 1) res = min(res, dat[l++]);
                if (r & 1) res = min(res, dat[--r]);
                l >>= 1; r >>= 1;
            }
            return res;
        }
    };
    vector<vector<int>> G;
    vector<int> in, ord;
    vector<int> dep;
    SegmentTree seg;

    void dfs(int v) {
        in[v] = ord.size();
        ord.push_back(v);
        for (int u : G[v]) {
            dep[u] = dep[v] + 1;
            dfs(u);
            ord.push_back(v);
        }
    }

public:
    void init(vector<vector<int>> G_, int r) {
        G = move(G_);
        int N = G.size();
        dep.assign(N, 0);
        in.assign(N, -1);
        ord.reserve(2 * N);
        dfs(r);
        vector<pair<int, int>> v;
        for (int i : ord) v.emplace_back(dep[i], i);
        seg.init(v);
    }
    int lca(int a, int b) const {
        int l = in[a], r = in[b];
        if (l > r) swap(l, r);
        return seg.prod(l, r + 1).second;
    }
    int depth(int v) const { return dep[v]; }
};

class WeightedUnionFind {
    vector<int> par;
    vector<ll> wei;

public:
    void init(int n) {
        par.resize(n, -1);
        wei.resize(n, 0);
    }
    int find(int v) {
        if (par[v] < 0) return v;
        int r = find(par[v]);
        wei[v] += wei[par[v]];
        return par[v] = r;
    }
    ll weight(int v) {
        find(v);
        return wei[v];
    }
    void merge(int u, int v, ll w) {
        w += weight(u);
        w -= weight(v);
        u = find(u);
        v = find(v);
        if (u == v) {
            assert(w == 0);
            return;
        }
        if (par[u] > par[v]) {
            swap(u, v);
            w = -w;
        }
        par[u] += par[v];
        par[v] = u;
        wei[v] = w;
    }
    bool same(int u, int v) { return find(u) == find(v); }
};

int N, K, Q, T;
vector<int> P1, P2, ord;
vector<vector<int>> G1, G2, Gaux;
LowestCommonAncestor lca1, lca2;
WeightedUnionFind uf1, uf2;

int main() {
    cin >> N >> K >> Q >> T;
    int r1 = -1, r2 = -1;
    P1.resize(N); P2.resize(N);
    G1.resize(N); G2.resize(N);
    rep (i, N) {
        cin >> P1[i];
        if (P1[i] == -1) r1 = i;
        else G1[--P1[i]].push_back(i);
    }
    rep (i, N) {
        cin >> P2[i];
        if (P2[i] == -1) r2 = i;
        else G2[--P2[i]].push_back(i);
    }
    lca1.init(G1, r1);
    lca2.init(G2, r2);
    vector<int> verts;
    {
        verts.reserve(K);
        auto dfs = [&](auto&& self, int v) -> void {
            if ((int)verts.size() == K) return;
            verts.push_back(v);
            for (int u : G1[v]) self(self, u);
        };
        dfs(dfs, r1);
    }
    assert((int)verts.size() == K);
    int raux;
    vector<int> aux;
    Gaux.resize(N);
    {
        ord.resize(N);
        {
            int cnt = 0;
            auto dfs = [&](auto&& self, int v) -> void {
                ord[v] = cnt++;
                for (int u : G2[v]) self(self, u);
            };
            dfs(dfs, r2);
        }
        sort(all(verts), [&](int a, int b) { return ord[a] < ord[b]; });
        aux.reserve(2 * K - 1);
        for (auto i : verts) {
            aux.push_back(i);
        }
        rep (i, K - 1) {
            aux.push_back(lca2.lca(verts[i], verts[i + 1]));
        }
        sort(all(aux), [&](int a, int b) { return ord[a] < ord[b]; });
        aux.erase(unique(all(aux)), aux.end());
        raux = aux.front();
        stack<int> st;
        for (int v : aux) {
            if (st.empty()) {
                st.push(v);
                continue;
            }
            while (lca2.lca(st.top(), v) != st.top()) {
                st.pop();
                assert(!st.empty());
            }
            Gaux[st.top()].push_back(v);
            st.push(v);
        }
    }
    vector<bool> is_aux(N, false);
    for (auto v : verts) is_aux[v] = true;

    sort(all(verts));
    rep (i, K) cout << verts[i] + 1 << " \n"[i == K - 1];
    vector<pair<int, int>> qs;
    {
        qs.reserve(K - 1);
        auto dfs = [&](auto&& self, int v) -> int {
            int res = -1;
            if (is_aux[v]) res = v;
            for (int u : Gaux[v]) {
                int r = self(self, u);
                if (res == -1) res = r;
                else {
                    qs.emplace_back(res, r);
                }
            }
            return res;
        };
        dfs(dfs, raux);
    }
    assert((int)qs.size() == K - 1);
    for (auto [a, b] : qs) cout << "? " << a + 1 << " " << b + 1 << "\n";
    cout << "!\n";
    cout.flush();

    uf1.init(N); uf2.init(N);
    for (auto [a, b] : qs) {
        ll d1a, d1b, d2a, d2b; cin >> d1a >> d1b >> d2a >> d2b;
        int l1 = lca1.lca(a, b), l2 = lca2.lca(a, b);
        uf1.merge(l1, a, d1a); uf1.merge(l1, b, d1b);
        uf2.merge(l2, a, d2a); uf2.merge(l2, b, d2b);
    }
    vector<pair<ll, ll>> ans;
    rep (_, T) {
        int a, b; cin >> a >> b;
        --a, --b;
        int l1 = lca1.lca(a, b), l2 = lca2.lca(a, b);
        ll d1 = uf1.weight(a) + uf1.weight(b) - 2 * uf1.weight(l1);
        ll d2 = uf2.weight(a) + uf2.weight(b) - 2 * uf2.weight(l2);
        ans.emplace_back(d1, d2);
    }
    assert(ans.size() == T);
    rep (i, ans.size()) {
        cout << ans[i].first << " " << ans[i].second << "\n";
    }
    cout.flush();
}

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 Main.cpp:1:
Main.cpp: In function 'int main()':
Main.cpp:272:23: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  272 |     assert(ans.size() == T);
      |            ~~~~~~~~~~~^~~~

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