Skip to content

Instantly share code, notes, and snippets.

@brunomaletta
Last active January 13, 2023 20:54
Show Gist options
  • Save brunomaletta/81ba61f5e1e508e5f7e1d536dfd19958 to your computer and use it in GitHub Desktop.
Save brunomaletta/81ba61f5e1e508e5f7e1d536dfd19958 to your computer and use it in GitHub Desktop.
#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);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment