Fusión small-to-large

Fusionar estructuras de datos

Es evidente que las listas enlazadas  se pueden fusionar en tiempo $\mathcal{O}(1)$. Pero ¿qué hay de los conjuntos o los vectores?

Consideremos un árbol enraizado en el nodo $1$, donde cada nodo tiene un color.

Para cada nodo, almacenemos un conjunto que contiene solo ese nodo, y queremos fusionar los conjuntos en el subárbol del nodo de modo que cada nodo tenga un conjunto formado por todos los colores del subárbol del nodo. Hacer esto nos permite resolver una variedad de problemas, como consultar la cantidad de colores distintos en cada subárbol.

Solución naive

Supongamos que queremos fusionar dos conjuntos $a$ y $b$ de tamaños $n$ y $m$, respectivamente. Una posibilidad es la siguiente:

for (int x : b) a.insert(x);

que corre en tiempo $\mathcal{O}(m\log (n+m))$, lo que da un tiempo de ejecución de $\mathcal{O}(N^2\log N)$ en el peor caso. Si en cambio mantenemos $a$ y $b$ como vectores ordenados, podemos fusionarlos en tiempo $\mathcal{O}(n+m)$, pero $\mathcal{O}(N^2)$ también es demasiado lento.

Mejor solución

Con solo una línea extra de código, podemos acelerar esto de forma significativa.

if (a.size() < b.size()) swap(a, b);
for (int x : b) a.insert(x);

Observar que swap  intercambia dos conjuntos en tiempo $\mathcal{O}(1)$. Así, fusionar un conjunto más pequeño de tamaño $m$ en el más grande de tamaño $n$ toma tiempo $\mathcal{O}(m\log n)$.

Afirmación: La solución corre en tiempo $\mathcal{O}(N\log^2N)$.

Demostración: Al fusionar dos conjuntos, se mueve del conjunto más pequeño al más grande. Si el tamaño del conjunto más pequeño es $X$, entonces el tamaño del conjunto resultante es al menos $2X$. Así, un elemento que se ha movido $Y$ veces estará en un conjunto de tamaño al menos $2^Y$, y como el tamaño máximo de un conjunto es $N$ (la raíz), cada elemento se moverá a lo sumo $\mathcal{O}(\log N$) veces.

Código completo

#include <bits/stdc++.h>
 
using namespace std;
 
const int MAX_N = 2e5;
 
// nodes will be 1-indexed like in the problem
vector<int> adj[MAX_N + 1];
 
set<int> colors[MAX_N + 1];
int distinct_num[MAX_N + 1];
 
void process_colors(int curr, int parent) {
	for (int n : adj[curr]) {
		if (n != parent) {
			process_colors(n, curr);
			// make x the larger set always
			if (colors[curr].size() < colors[n].size()) {
				swap(colors[curr], colors[n]);
			}
			for (int item : colors[n]) { colors[curr].insert(item); }
		}
	}
	distinct_num[curr] = colors[curr].size();
}
 
int main() {
	int n;
	cin >> n;
	for (int i = 1; i <= n; i++) {
		int a;
		cin >> a;
		colors[i].insert(a);
	}
	for (int i = 1; i < n; i++) {
		int a;
		int b;
		cin >> a >> b;
		adj[a].push_back(b);
		adj[b].push_back(a);
	}
	process_colors(1, 0);
	for (int i = 1; i <= n; i++) { cout << distinct_num[i] << (i < n ? " " : "\n"); }
}

Generalización

También podemos fusionar otras estructuras de datos de la librería estándar como std::map o std:unordered_map de la misma forma. Sin embargo, std::swap no siempre corre en tiempo $\mathcal{O}(1)$. Por ejemplo, intercambiar std::arrays toma tiempo lineal en la suma de los tamaños de los arreglos, y lo mismo ocurre con las estructuras de datos policy-based de GCC  como __gnu_pbds::tree o __gnu_pbds::gp_hash_table.

Para intercambiar dos estructuras de datos policy-based a y b en tiempo $\mathcal{O}(1)$, usar a.swap(b) en su lugar. Observar que para las estructuras de datos de la librería estándar, swap(a,b) es equivalente a a.swap(b).

Problemas