Идея

Заведём дополнительный массив rank. По iой позиции будет хранится ранг вершины x — расстояние от листа-наследника до x.

Что улучшит

Улучшить операцию union(x, y): лучше подвешивать невысокое дерево к высокому, чем наоборот. Тогда высота не изменится. Если же ранги равны, то только тогда высота изменится на 1.

После этой эвристики высота дерева увеличивается <=> ранги корней объединяемых деревьев равны.

Улучшенный union(x, y) $\approx O(1)$, см пояснение на странице ранее

Тут мы используем улучшенный find(x), который использует эвристику сжатия путей (Path compression)

void union(x, y) {
	xRoot = find(x); 
	yRoot = find(y); 
	if (rank[xRoot] > rank[yRoot]) {
		arr[yRoot] = xRoot; // подвешиваем yRoot к xRoot
	} else {
		arr[xRoot] = yRoot; // подвешиваем xRoot к yRoot
		
		// если ранги равны, то высота дерева увеличится на 1
		if (rank[xRoot] == rank[yRoot]) { 
			rank[yRoot] += 1; 
		}
	}
}