Tropical Garden

Explicación

Primero, en lugar de atacar el problema para valores arbitrarios de $k$ , calculemos cuántos pasos le toma a cada nodo llegar al nodo $p$ .

En una fuente, solo tomaremos siempre los dos senderos salientes más hermosos. Consideremos un grafo en el que dibujamos aristas entre estados. Cada estado se representa como un par $(u, f)$ , donde:

$u$ es nuestro nodo actual
$f$ indica si estamos tomando el sendero más hermoso
- $f=0$ significa que podemos tomar el sendero más hermoso
- $f=1$ significa que solo podemos tomar el segundo sendero más hermoso

Ahora queremos hallar la distancia desde todo otro estado hasta $(p, 0)$ . Hallar las distancias hasta $(p, 1)$ se puede manejar de forma similar.

Podemos usar el método de recorrido de grafos que prefiramos para recorrer desde $(p, 0)$ hacia todos los demás estados. Para esto construimos el grafo invertido y recorremos desde $(p, 0)$ hacia los otros estados. Tras calcular las distancias, sabemos que la distancia desde cualquier nodo $i$ hasta nuestro estado $(p, 0)$ es la distancia desde $(i, 0)$ hasta $(p, 0)$ . Tratamos $(p, 1)$ de forma similar.

Bien, ahora sabemos cómo calcular la distancia desde cada nodo hasta $p$ . ¿Cómo manejamos valores grandes de $k$ ?

Consideremos que cada estado se mapea de forma directa a otro estado, formando un grafo de sucesores (successor graph). Así, si en algún momento llegamos a $(p, 0)$ o $(p, 1)$ , debe ocurrir una de las siguientes:

Llegamos a $p$ y terminamos en un ciclo donde $p$ no está presente
Llegamos a $(p, f)$ y ciclamos de vuelta a $(p, f)$
Llegamos a $(p, f)$ , encontramos $(p, 1-f)$ y volvemos a $(p, f)$

Supongamos que nuestro estado $(p, f)$ cae en el segundo escenario, donde termina en un ciclo de longitud $n$ . Entonces, para que un nodo $i$ pueda alcanzar $(p, f)$ después de $k$ senderos, debe cumplirse lo siguiente para algún entero no negativo $c$ :

k = \text{dist}((p, f), (i, 0)) + c \cdot n.

Los escenarios 1 y 3 se pueden manejar de forma similar. Para el escenario 3, hay que guardar el instante en que visitamos el estado $(p, 1 - f)$ en nuestro ciclo.

Con esto, podemos resolver cada consulta en $\mathcal{O}(N)$ comprobando cada nodo y viendo si termina en $p$ , usando los tres casos descritos arriba.

Implementación

Complejidad temporal: $\mathcal{O}(M + NQ)$

C++


#include "garden.h"
#include "gardenlib.h"
#include <bits/stdc++.h>
 
using ll = long long;
 
constexpr int INF = 1e9;
 
void count_routes(int n, int m, int p, int r[][2], int q, int g[]) {
	std::vector<std::vector<std::array<int, 2>>> adj(n);
	for (int i = 0; i < m; i++) {
		adj[r[i][0]].push_back({r[i][1], i});
		adj[r[i][1]].push_back({r[i][0], i});
	}
 
	std::vector<std::vector<std::array<int, 2>>> rev(n);
	for (int i = 0; i < n; i++) {
		// solo nos importan las dos mejores aristas
		// las aristas ya están en orden
		if (adj[i].size() > 2) adj[i].resize(2);
		for (const auto &[j, w] : adj[i]) { rev[j].push_back({i, w}); }
	}
 
	std::vector<std::array<int, 2>> dist(n, {INF, INF});
	std::vector<std::array<int, 2>> to_p(n, {INF, INF});
	for (int tt = 0; tt < 2; tt++) {
		// tt = 0 significa que tomamos la mejor arista desde p
		// tt = 1 significa que no tomamos la mejor arista
		if (tt == 1 && adj[p].size() == 1) break;
 
		std::queue<std::array<int, 3>> bfs;
		bfs.push({0, p, tt});
		while (!bfs.empty()) {
			const auto [t, u, f] = bfs.front();
			bfs.pop();
			if (dist[u][f] != INF) continue;
			dist[u][f] = t;
 
			int wt = adj[u][0][1];
			for (const auto &[j, w] : rev[u]) {
				// nos aseguramos de que podamos tomar la arista actual
				if (f == 0 && w == wt && adj[u].size() > 1) continue;
				if (f == 1 && w != wt && adj[u].size() > 1) continue;
 
				int flag = w != adj[j][0][1];
				bfs.push({t + 1, j, flag});
			}
		}
 
		// la distancia será dist[i][0] porque siempre empezamos con la mejor arista
		for (int i = 0; i < n; i++) { to_p[i][tt] = dist[i][0]; }
 
		dist.assign(n, {INF, INF});
	}
 
	/** @return {longitud del ciclo de (p, f), instante en que visitamos (p, 1 - f)} */
	const auto get_cycle = [&](int f) -> std::array<int, 2> {
		int cycle_len = 0;
		int saw = 0;
		std::vector<std::array<bool, 2>> vis(n);
 
		// seguimos visitando nodos hasta volver al nodo p o alcanzar un ciclo
		int node = p, flag = f, trav = 0;
		while (!vis[node][flag]) {
			vis[node][flag] = true;
			if (node == p && flag == !f) saw = trav;
 
			const auto [nxt, wt] = adj[node][flag];
			if (adj[nxt].size() == 1) {
				node = nxt, flag = 0;
			} else {
				node = nxt, flag = (adj[nxt][0][1] == wt);
			}
 
			trav++;
		}
 
		if (node == p) cycle_len = trav;
		else cycle_len = INT_MAX, saw = 0;
		return {cycle_len, saw};
	};
 
	std::array<int, 2> s1 = get_cycle(0);
	std::array<int, 2> s2 = {INT_MAX, 0};
	if (adj[p].size() > 1) s2 = get_cycle(1);
 
	for (int t = 0; t < q; t++) {
		int k = g[t], res = 0;
		for (int i = 0; i < n; i++) {
			// comprobamos to_p[i][0] y to_p[i][1] para ver si ciclan hacia p
			if (to_p[i][0] <= k &&
			    ((k - to_p[i][0]) % s1[0] == 0 || (k - to_p[i][0]) % s1[0] == s1[1])) {
				res++;
			} else if (to_p[i][1] <= k && ((k - to_p[i][1]) % s2[0] == 0 ||
			                               (k - to_p[i][1]) % s2[0] == s2[1])) {
				res++;
			}
		}
 
		answer(res);
	}
}