priority_queue_ds.cpp

#include <bits/stdc++.h>

using namespace std;

#define _ ios_base::sync_with_stdio(0);cin.tie(0);
#define endl '\n'

typedef long long ll;

const int INF = 0x3f3f3f3f;
const ll LINF = 0x3f3f3f3f3f3f3f3fll;

struct dsu {
	vector<int> id, sz;
	stack<stack<pair<int&, int>>> st;

	dsu(int n) : id(n), sz(n, 1) { 
		iota(id.begin(), id.end(), 0), st.emplace(); 
	}
		
	void save(int &x) { st.top().emplace(x, x); }

	void checkpoint() { st.emplace(); }
	
	void rollback() {
		while (st.top().size()) {
			auto [end, val] = st.top().top(); st.top().pop();
			end = val;
		}
		st.pop();
	}

	int find(int a) { return a == id[a] ? a : find(id[a]); }

	void unite(int a, int b) {
		a = find(a), b = find(b);
		if (a == b) return;
		if (sz[a] < sz[b]) swap(a, b);
		save(sz[a]), save(id[b]);
		sz[a] += sz[b], id[b] = a;
	}
	void update(pair<int, int> p) {
	    checkpoint();
	    unite(p.first, p.second);
	}
};

// assumes all priorities are distinct
template<typename DS, typename UPD> struct priority_queue_ds {
    DS D;
    vector<tuple<UPD, int, int>> upd; // {u, p, idx_in_pos}
    set<pair<int, int>> st;
    vector<int> pos;
    
    priority_queue_ds(int n) : D(n) {}
    
    void update(UPD u, int p) { // O(log n + T(n))
        D.update(u);
        st.emplace(p, pos.size());
        upd.emplace_back(u, p, pos.size());
        pos.push_back(upd.size() - 1);
    }
    
    int query(int a) { // O(T(n))
        return D.find(a);
    }
    
    void pop() { // O(log n + k T(n))
        int k = 1, min_p; // k = number of pops we will do
        vector<tuple<UPD, int, int>> small, big;
        auto it = st.end();
        for (int qt = 0; qt++ < (k+1)/2;) {
            it--;
            min_p = it->first;
            int i = pos[it->second];
            if (qt > 1) big.push_back(upd[i]);
            k = max<int>(k, upd.size() - i);
        }
        
        for (int i = 0; i < k; i++) {
            D.rollback();
            auto [u, p, idx] = upd.rbegin()[i];
            if (p < min_p) small.emplace_back(u, p, idx);
        }
        
        st.erase(prev(st.end()));
        upd.erase(upd.end() - k, upd.end());
        
        small.insert(small.end(), big.rbegin(), big.rend());
        for (auto [u, p, idx] : small) {
            D.update(u);
            upd.emplace_back(u, p, idx);
            pos[idx] = upd.size() - 1;
        }
    }
};

int main() { _
    priority_queue_ds<dsu, pair<int, int>> PQ(5);
    priority_queue<tuple<int, int, int>> pq;
    auto upd = [&](int a, int b, int p) {
      PQ.update({a, b}, p);
      pq.emplace(p, a, b);
    };
    auto print = [&]() {
        for (int i = 0; i < 5; i++) cout << i << " " << PQ.query(i) << endl;
    };
    auto pop = [&]() {
        PQ.pop();
        auto [p, a, b] = pq.top(); pq.pop();
        cout << "popped " << a << " " << b << endl;
    };
    upd(1, 2, 5);
    upd(0, 1, 6);
    upd(4, 1, 3);
    upd(1, 3, 1);
    pop();
    pop();
    print();
    exit(0);
}