hackercupai/hackercup huggingface dataset:
We can use a disjoint set union (DSU) data structure to maintain the colors of all stakes. We'll keep track of the color of a component by updating the color of that set's representation (a.k.a its root or parent), and also the inverse. The repaint function will look something like this:

```
void repaint(int c1, int c2) {
int p1 = parent_of_color[c1];
if (p1 == -1) {
return;
}
int p2 = parent_of_color[c2];
if (p2 == -1) {
color_of_parent[p1] = c2;
parent_of_color[c1] = -1;
parent_of_color[c2] = p1;
return;
}
unite_sets(p1, p2);
color_of_parent[p1] = color_of_parent[p2] = c2;
parent_of_color[c1] = -1;
parent_of_color[c2] = parent[p2];
}
```
We also need to set \(\text{color\_of\_parent}[i] = \text{color\_of\_parent}[j]\) when updating \(\text{parent}[i] = j\) elsewhere in the DSU.

For this problem, there are \(B = \lceil N/K \rceil\) different "bins", for which we'd like to see if the colors are uniform for all of the bins after each repainting.

One possible idea is to maintain \(B\) hash tables \(H_{1..B}\), one for each bin, where \(H_{b}[r]\) stores the number of balls with DSU representative \(r\) in bin \(b\). If \(H_b[r]\) drops to zero, we delete the key \(r\) it so that the number of keys in \(H_b\) is alway equal to the number of distinct representatives in bin \(b\). Note that if there is only \(1\) key in \(H_b\), then bin \(b\) is uniformly colored. We increment \(H_{\text{bin}(i)}[j]\) whenever we update \(\text{parent}[i] = j\) in the DSU (and decrement the old parent's count).

We can also maintain another hash set \(S\) of the bin numbers which are currently uniform. Check if the size of \(H_b\) changed from \(2\) to \(1\) after every update (if so, add \(b\) to \(S\)) or if the size \(H_b\) changed from \(1\) to \(2\) (then delete \(b\) from S). The first time that \(|S| = B\) is the answer.

The only catch is that using the classic path compression DSU implementation, the \(\text{parent}[]\) array is lazily updated when \(\text{find\_parent}()\) is called (and might not be consistent for all nodes after every \(\text{unite\_sets}()\) call). Instead, we can [store the DSU explicitly in a set list](https://cp-algorithms.com/data_structures/disjoint_set_union.html#storing-the-dsu-explicitly-in-a-set-list-applications-of-this-idea-when-merging-various-data-structures), which guarantees consistency after each union. The running time across \(N\) elements and \(M\) union operations is \(\mathcal{O}(M + N \log N)\) for this variant of DSU, where \(\mathcal{O}(N \log N)\) is the max number times an element will have to be transferred from one set to another. Each such move requires updating \(H\) and \(S\) in \(\mathcal{O}(1)\) time. Thus the overall time complexity is \(\mathcal{O}(M + N \log N)\).