Distinct Colors

Solución 1 - Fusión small-to-large

Solución 2 - PURS

Consideremos el tour de Euler del árbol. Aplanamos el árbol en un arreglo, donde cada nodo corresponde a un rango de este arreglo. Ahora, la tarea es esencialmente hallar el número de valores distintos en un rango $N$ veces (una vez por cada nodo).

Consideremos iterar el arreglo del tour de Euler de izquierda a derecha. Cuando consideramos un nodo en el índice $i$ del arreglo del tour de Euler, el rango de su subárbol será $[j, i]$ , donde $j\leq i$ . Ahora, enfoquémonos en un solo color. Observemos que si hay múltiples colores $c$ a la izquierda de cierto índice $i$ , solo la ocurrencia más a la derecha de $c$ es relevante. Solo queremos contar $c$ una vez, así que elegimos contar solo la ocurrencia más a la derecha. Esto es porque cualquier ocurrencia de $c$ que no sea la más a la derecha desde $i$ debe incluir la ocurrencia más a la derecha de $c$ . Más formalmente, cualquier segmento $[l, r]$ debe contener $i$ si $l <= i <= r$ .

Con esta observación, podemos reducir esencialmente las consultas de colores distintos a una consulta de rango simple. A medida que iteramos el tour de Euler de izquierda a derecha, marcamos el índice actual (en el BIT) como $1$ , y si el color del nodo en el índice actual ya ocurrió antes, lo ponemos en $0$ en el BIT. Podemos hallar la solución de cada nodo haciendo una consulta de suma mientras iteramos.

Implementación

Complejidad temporal: $\mathcal O(N\log N)$

C++


#include <bits/stdc++.h>
 
using namespace std;
#define pb push_back
 
 
struct BIT {
	vector<int> bit;
	int n;
	BIT(int n) : n(n + 1), bit(n + 1) {}
	int sum(int r) {
		r++;
		int ret = 0;
		while (r > 0) {
			ret += bit[r];
			r -= r & -r;
		}
		return ret;
	}
 
	void update(int idx, int v) {
		idx++;
		while (idx < n) {
			bit[idx] += v;
			idx += idx & -idx;
		}
	}
};
 
const int MAXN = 2e5 + 1;
 
/*
 * tour = tour de Euler
 * color = color de cada nodo
 * answer = respuesta de cada nodo
 * lend = extremo izquierdo de cada nodo (como se explica en el editorial)
 */
int tour[MAXN], color[MAXN], answer[MAXN], lend[MAXN];
 
vector<int> adj[MAXN];
int idx = 0;
 
void dfs(int u, int par = 0) {
	lend[u] = idx;
	for (int n : adj[u]) {
		if (n == par) continue;
		dfs(n, u);
	}
	tour[idx] = u;
	idx++;
}
 
int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	int N, u, v;
	cin >> N;
	for (int i = 1; i <= N; i++) { cin >> color[i]; }
	for (int i = 0; i < N - 1; i++) {
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	dfs(1);
	BIT bit(N);
 
	// Este mapa guarda la ocurrencia más a la derecha de cada color
	map<int, int> last;
	for (int i = 0; i < N; i++) {
		// Si alguna vez consideramos el color de tour[i]
		if (last.count(color[tour[i]])) {
			// La última vez que apareció el color de tour[i], lo volvemos a poner en 0
			bit.update(last[color[tour[i]]], -1);
		}
		// Cambiamos la última ocurrencia del color de tour[i]
		last[color[tour[i]]] = i;
		// Reflejamos ese cambio en nuestro BIT
		bit.update(i, 1);
		/*
		 * La respuesta del nodo de tour[i] es simplemente
		 * la suma de 1s en su subintervalo contiguo del tour de Euler
		 */
		answer[tour[i]] = bit.sum(i) - bit.sum(lend[tour[i]] - 1);
	}
	for (int i = 1; i <= N; i++) { cout << answer[i] << " "; }
}