Submission #1011282

#	Time	Username	Problem	Language	Result	Execution time	Memory
1011282		Forested	Prize (CEOI22_prize)	C++17	35 / 100	2193 ms	478844 KiB

Main

#include <bits/stdc++.h>
using namespace std;
using i32 = int;
using i64 = long long;
template <typename T>
using V = vector<T>;
template <typename T>
using VV = V<V<T>>;
#define OVERRIDE4(a, b, c, d, ...) d
#define REP2(i, n) for (i32 i = 0; i < (i32)(n); ++i)
#define REP3(i, l, r) for (i32 i = (i32)(l); i < (i32)(r); ++i)
#define REP(...) OVERRIDE4(__VA_ARGS__, REP3, REP2)(__VA_ARGS__)
#define PER2(i, n) for (i32 i = (i32)(n)-1; i >= 0; --i)
#define PER3(i, l, r) for (i32 i = (i32)(r)-1; i >= (i32)(l); --i)
#define PER(...) OVERRIDE4(__VA_ARGS__, PER3, PER2)(__VA_ARGS__)
#define LEN(x) (i32) size(x)
#define ALL(x) begin(x), end(x)
template <typename T>
bool chmin(T &x, const T &y) {
    if (x > y) {
        x = y;
        return true;
    }
    return false;
}
template <typename T>
bool chmax(T &x, const T &y) {
    if (x < y) {
        x = y;
        return true;
    }
    return false;
}

struct HLD {
    VV<i32> g;
    V<i32> par, sz, dep, head, in, out;

    void dfs0(i32 v) {
        if (!g[v].empty() && g[v].front() == par[v]) {
            swap(g[v].front(), g[v].back());
        }
        for (i32 &u : g[v]) {
            if (u == par[v]) {
                continue;
            }
            par[u] = v;
            dep[u] = dep[v] + 1;
            dfs0(u);
            sz[v] += sz[u];
            if (sz[u] > sz[g[v].front()]) {
                swap(g[v].front(), u);
            }
        }
    }

    void dfs1(i32 v, i32 &t) {
        in[v] = t++;
        for (i32 u : g[v]) {
            if (u == par[v]) {
                continue;
            }
            head[u] = (u == g[v].front() ? head[v] : u);
            dfs1(u, t);
        }
        out[v] = t;
    }

    HLD(const VV<i32> &g, i32 root)
        : g(g),
          par(LEN(g), -1),
          sz(LEN(g), 1),
          dep(LEN(g), 0),
          head(LEN(g), root),
          in(LEN(g), 0),
          out(LEN(g), 0) {
        dfs0(root);
        i32 t = 0;
        dfs1(root, t);
    }

    i32 lca(i32 u, i32 v) const {
        while (u != v) {
            if (in[u] > in[v]) {
                swap(u, v);
            }
            if (head[u] == head[v]) {
                break;
            }
            v = par[head[v]];
        }
        return u;
    }
};

class UnionFindWithPotential {
    V<i32> par;
    V<i32> diff;

    i32 root(i32 v) {
        if (par[v] < 0) {
            return v;
        } else {
            i32 p = par[v];
            i32 r = root(par[v]);
            diff[v] += diff[p];
            par[v] = r;
            return r;
        }
    }

public:
    UnionFindWithPotential(i32 n) : par(n, -1), diff(n, 0) {}

    // a[u] - a[v] = x
    void unite(i32 u, i32 v, i32 x) {
        i32 ru = root(u), rv = root(v);
        if (ru == rv) {
            return;
        }
        if (-par[ru] > -par[rv]) {
            swap(ru, rv);
            swap(u, v);
            x = -x;
        }
        par[rv] += par[ru];
        par[ru] = rv;
        diff[ru] = x - diff[u] + diff[v];
    }

    // a[u] - a[v]
    i32 getd(i32 u, i32 v) {
        i32 r = root(u);
        i32 r2 = root(v);
        assert(r == r2);
        return diff[u] - diff[v];
    }
    
    void show() {
        cout << '\n';
        REP(i, LEN(par)) {
            i32 r = root(i);
            cout << r << ' ' << diff[i] << '\n';
        }
        cout << '\n';
    }
};

int main() {
    i32 n, k, q, t;
    cin >> n >> k >> q >> t;
    V<i32> t0(n), t1(n);
    REP(i, n) {
        cin >> t0[i];
        if (t0[i] > 0) {
            --t0[i];
        }
    }
    REP(i, n) {
        cin >> t1[i];
        if (t1[i] > 0) {
            --t1[i];
        }
    }
    i32 r0 = (i32)(find(ALL(t0), -1) - t0.begin());
    i32 r1 = (i32)(find(ALL(t1), -1) - t1.begin());
    VV<i32> g0(n), g1(n);
    REP(i, n) {
        if (t0[i] != -1) {
            g0[t0[i]].push_back(i);
        }
    }
    REP(i, n) {
        if (t1[i] != -1) {
            g1[t1[i]].push_back(i);
        }
    }
    HLD hld0(g0, r0), hld1(g1, r1);

    
    V<i32> st(n);
    iota(ALL(st), 0);
    {
        mt19937 mt(998);
        shuffle(ALL(st), mt);
    }
    st.resize(k);
    V<i32> ord0 = st;
    V<i32> ord1 = st;
    sort(ALL(ord0),
         [&](i32 u, i32 v) -> bool { return hld0.in[u] < hld0.in[v]; });
    sort(ALL(ord1),
         [&](i32 u, i32 v) -> bool { return hld1.in[u] < hld1.in[v]; });
    V<pair<i32, i32>> qrys;
    if (q == k - 1) {
        REP(i, k - 1) { qrys.emplace_back(ord0[i], ord0[i + 1]); }
    } else {
        REP(i, k - 1) { qrys.emplace_back(ord0[i], ord0[i + 1]); }
        REP(i, k - 1) { qrys.emplace_back(ord1[i], ord1[i + 1]); }
    }
    REP(i, k) { cout << st[i] + 1 << " \n"[i + 1 == k]; }
    for (auto [a, b] : qrys) {
        cout << "? " << a + 1 << ' ' << b + 1 << '\n';
    }
    cout << "!\n";
    cout.flush();
    using PCT = tuple<i32, i32, i32, i32>;
    V<PCT> resp(q);
    for (auto &[x, y, z, w] : resp) {
        cin >> x >> y >> z >> w;
    }
    V<pair<i32, i32>> dq(t);
    REP(i, t) {
        cin >> dq[i].first >> dq[i].second;
        --dq[i].first;
        --dq[i].second;
    }

    UnionFindWithPotential uf0(n), uf1(n);
    if (q == k - 1) {
        REP(i, k - 1) {
            auto [a, b] = qrys[i];
            i32 lca0 = hld0.lca(a, b);
            i32 lca1 = hld1.lca(a, b);
            auto [x, y, z, w] = resp[i];
            uf0.unite(a, lca0, x);
            uf0.unite(b, lca0, y);
            uf1.unite(a, lca1, z);
            uf1.unite(b, lca1, w);
        }
    } else {
        REP(i, k - 1) {
            auto [a, b] = qrys[i];
            i32 lca0 = hld0.lca(a, b);
            auto [x, y, z, w] = resp[i];
            uf0.unite(a, lca0, x);
            uf0.unite(b, lca0, y);
        }
        REP(i, k - 1) {
            auto [a, b] = qrys[k - 1 + i];
            i32 lca1 = hld1.lca(a, b);
            auto [x, y, z, w] = resp[k - 1 + i];
            uf1.unite(a, lca1, z);
            uf1.unite(b, lca1, w);
            //uf1.show();
        }
    }
    
    V<pair<i32, i32>> dans(t);
    REP(i, t) {
        auto [a, b] = dq[i];
        i32 lca0 = hld0.lca(a, b);
        i32 lca1 = hld1.lca(a, b);
        dans[i].first = uf0.getd(a, lca0) + uf0.getd(b, lca0);
        dans[i].second = uf1.getd(a, lca1) + uf1.getd(b, lca1);
    }
    REP(i, t) { cout << dans[i].first << ' ' << dans[i].second << '\n'; }
    cout.flush();
}

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