GCD Harmony

Hagamos que cada nodo contenga un arreglo de DP con el menor costo de hacer armonioso el subárbol de ese nodo si el nodo mismo tuviera que ser un cierto valor.

Si se obtiene WA con el estado de DP de arriba, consideremos un árbol con dos ramas largas, una con todos los valores iguales a 97, y otra con todos los valores iguales a 89. La raíz es 1.

¿Funcionaría el programa en un caso así?

Si solo tomamos valores de 2 a 100, la DP no funcionará: tiene que contemplar valores mucho mayores que eso. ¿Cómo podríamos hacer esto sin exceder el límite de tiempo?

Explicación

DP inicial

Mantengamos un arreglo de DP (llamado $dp$) para cada nodo con la definición dada en las pistas.

$dp$ empieza con el valor inicial de solo cambiar el nodo a cada valor de partida. Por ejemplo, si el valor inicial de un nodo era $5$, entonces $dp[5]=0$, mientras que $dp[7]=7$.

Para completar este arreglo, iteramos recursivamente por todos los hijos y luego comprobamos cuál es su valor de menor costo que sigue siendo compatible con cada elemento del $dp$ del nodo actual. Por ejemplo, si consideramos $dp[3]$, comprobaríamos el $dp[3]$, $dp[6]$, etc. de los hijos.

Refinamiento

Como se indicó en las pistas, solo guardar valores de 2 a 100 no funciona. El costo máximo de hacer armonioso un árbol es $2 \cdot 5000=10^4$, así que tenemos que guardar valores hasta ese.

¿Cómo transicionamos, entonces, con estados tan grandes?

Notemos primero que nunca tendría sentido que un nodo tenga un valor final con un factor primo mayor que 100. Esto es cierto porque ¿por qué querríamos cambiar un nodo a un valor que tenga, por ejemplo, 101 como factor primo? Como el valor inicial de un nodo nunca puede exceder 100, este factor extra agrega costo adicional sin permitir que el nodo se armonice con ningún otro nodo adyacente.

Luego, para cada nodo, podemos llevar la cuenta del costo mínimo asociado a cada factor primo además de $dp$. Por ejemplo, el valor $7$ en este mapa llevaría la cuenta del costo mínimo entre $dp[7]$, $dp[14]$, etc. Podemos calcular este mapa después de que $dp$ se haya completado, así que no tomará demasiado tiempo.

Luego, al completar $dp$, en lugar de recorrer todos los $10^4$ valores posibles del hijo, miramos todos los factores primos del valor del nodo actual junto con el mapa auxiliar del hijo para calcular el mejor costo del hijo de forma eficiente.

Para ver por qué esto ahorra tiempo, digamos que estamos calculando $dp[51]$ para un cierto nodo. Si nos quedáramos con la transición vieja, tendríamos que iterar por los $10^4$ valores de cada hijo y comprobar si eran armoniosos con $51$. Con esta nueva transición, podemos obtener nuestra respuesta haciendo dos consultas al mapa por cada hijo: una para el $3$ y otra para el $17$. ¡Esto es mucho más eficiente que el enfoque anterior!

Implementación

Complejidad temporal: $\mathcal{O}(n^2 \pi(v))$, donde $\pi(v)$ denota la cantidad de primos menores o iguales que $v$.

#include <algorithm>
#include <cstdint>
#include <iostream>
#include <map>
#include <vector>
 
using std::cout;
using std::endl;
using std::vector;
 
const int MAX_VAL = 100;       // el mayor valor que puede tener un nodo
const int MAX_POSS = 2 * 5e3;  // definido como 2 * (máx. # de nodos)
 
const vector<int> PRIMES = {2,  3,  5,  7,  11, 13, 17, 19, 23, 29, 31, 37, 41,
                            43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97};
 
class PrimeLog {
  private:
	vector<vector<int>> rel_factors;
 
  public:
	PrimeLog(int n) : rel_factors(n + 1) {
		for (int i = 2; i <= n; i++) {
			for (int p : PRIMES) {
				if (i % p == 0) { rel_factors[i].push_back(p); }
			}
		}
	}
 
	/** @return the relevant prime factors of n */
	const vector<int> &get_factors(int n) { return rel_factors[n]; }
} factors(MAX_POSS);
 
class GCDTree {
  private:
	const vector<vector<int>> &neighbors;
	const vector<int> &vals;
	vector<vector<int>> min_cost;            // contiene el arreglo de dp de cada nodo
	vector<std::map<int, int>> factor_best;  // el mapa discutido arriba
 
	void dfs(int at, int prev) {
		min_cost[at] = vector<int>(MAX_POSS + 1);  // los índices 0 y 1 no se usan
		for (int i = 2; i <= MAX_POSS; i++) {
			if (i != vals[at]) { min_cost[at][i] = i; }
			// los primos mayores que 100 son irrelevantes
			if (factors.get_factors(i).empty()) { min_cost[at][i] = INT32_MAX; }
		}
 
		for (int n : neighbors[at]) {
			if (n == prev) { continue; }
 
			dfs(n, at);
			for (int i = 2; i <= MAX_POSS; i++) {
				if (factors.get_factors(i).empty()) { continue; }
				int n_min_cost = INT32_MAX;
				// con el enfoque viejo, este bucle también tomaría MAX_POSS
				// iteraciones, ¡demasiado lento!
				for (int p : factors.get_factors(i)) {
					n_min_cost = std::min(n_min_cost, factor_best[n][p]);
				}
				min_cost[at][i] += n_min_cost;
			}
		}
 
		// llenamos nuestro mapa ahora que el arreglo está completo
		for (int p : PRIMES) {
			int p_best = INT32_MAX;
			for (int i = p; i <= MAX_POSS; i += p) {
				p_best = std::min(p_best, min_cost[at][i]);
			}
			factor_best[at][p] = p_best;
		}
	}
 
  public:
	GCDTree(const vector<vector<int>> &neighbors, const vector<int> &vals)
	    : neighbors(neighbors), vals(vals), min_cost(neighbors.size()),
	      factor_best(neighbors.size()) {
		dfs(0, 0);
	}
 
	int min_adjustment() const {
		int res = INT32_MAX;
		// fuerza bruta de todos los valores posibles para la raíz
		for (int i = 2; i <= MAX_POSS; i++) { res = std::min(res, min_cost[0][i]); }
		return res;
	}
};
 
int main() {
	int node_num;
	std::cin >> node_num;
	vector<int> vals(node_num);
	for (int &v : vals) { std::cin >> v; }
	vector<vector<int>> neighbors(node_num);
	for (int i = 0; i < node_num - 1; i++) {
		int a, b;
		std::cin >> a >> b;
		neighbors[--a].push_back(--b);
		neighbors[b].push_back(a);
	}
 
	cout << GCDTree(neighbors, vals).min_adjustment() << endl;
}