dsu_rollback.cpp

// Noic - Ideia 5
// Exemplo 3

#include <bits/stdc++.h>

using namespace std;
const int maxn = 100010;

int pai[maxn], h[maxn]; // Declaro o pai e a altura do Union Find

struct updt { // Estrutura que guarda uma modificação
	int x, y; // x representa o pai e y o vértice modificado
	bool r; // Guarda se a altura foi alterada
};
stack<updt> st; // Pilha de modificações

int find(int x) {
	if(pai[x] == x)
		return x;

	return find(pai[x]);
}

void join(int x, int y) {
	x = find(x);
	y = find(y);

	if (h[x] < h[y])
		swap(x, y); // Faremos com que x guarde a componente de maior altura

	updt u;
	u.x = x; // Guarda o pai
	u.y = y; // Guarda o filho
	u.r = (h[x] == h[y]); // Guarda se a altura foi alterada
	st.push(u); // Insiro a modificação na pilha

	pai[y] = x;
	if (h[x] == h[y])
		h[x]++; // Se a altura for igual, incremetar a altura de x
}

void rollback() {
	if(st.empty())
		return; // Retornar a função caso não exista modificações

	updt u = st.top(); // Pegamos a modificação no topo da pilha
	st.pop(); // Removemos o topo da pilha

	pai[u.y] = 0; // Fazendo o pai de y ser nulo novamente
	if(u.r == true)
		h[u.x]--; // Se a altura for alterada, diminuir a altura;
}