Root LCA Queries

Pista 1

Sean $a$ , $b$ y $c$ los tres nodos dados en cada consulta. ¿Cuándo el número de raíces posibles será mayor que cero?

Respuesta a la pista 1

La respuesta de una consulta dada solo es mayor que cero cuando $c$ está en el camino de $a$ a $b$ , ya que por definición el LCA debe yacer en el camino de $a$ a $b$ . Ahora hay que hacer un poco de análisis por casos para calcular la respuesta.

Solución

Explicación

Como se señaló en las pistas, el LCA de $a$ y $b$ debe yacer en el camino entre estos nodos. Enraizamos el árbol de forma arbitraria y luego consideramos los siguientes casos:

Caso 1: $c$ no está en el camino de $a$ a $b$ .

La respuesta es $0$ para este caso.

Caso 2: $c$ es igual al LCA de $a$ y $b$ , para nuestra raíz elegida.

Cualquier nodo fuera del subárbol de $c$ es una raíz adecuada. Dentro del subárbol de $c$ , cualquier nodo cuyo camino al nodo $c$ no intersecte el camino de $a$ a $b$ es una raíz adecuada.

Caso 3: $c$ está en el camino de $a$ a $b$ , pero no es igual al LCA de $a$ y $b$ .

Este caso se maneja de la misma forma que el caso 2, excepto que solo los nodos en el subárbol de $c$ pueden ser raíces posibles.

Para responder consultas rápido, usamos binary lifting para calcular el LCA, y precalculamos el número de nodos en cada subárbol para quitar rápidamente los subárboles no adecuados de nuestra respuesta.

Implementación

Complejidad temporal: $\mathcal{O}(N\log{N} + Q\log{N})$

C++


#include <bits/stdc++.h>
using namespace std;
 
class Tree {
  private:
	const int n, log2dist;
	const vector<vector<int>> &adj;
	vector<vector<int>> lift;
	vector<int> dep, sub;
 
	/** calculates lift, dep, and sub arrays */
	void dfs(int u, int p) {
		for (int v : adj[u]) {
			if (v == p) continue;
			dep[v] = dep[u] + 1;
			lift[0][v] = u;
			for (int p = 1; p < log2dist; p++) {
				lift[p][v] = lift[p - 1][lift[p - 1][v]];
			}
			dfs(v, u);
			sub[u] += sub[v];
		}
	}
 
  public:
	Tree(vector<vector<int>> &_adj)
	    : n((int)_adj.size()), log2dist(1 + (int)log2(n)), adj(_adj),
	      lift(log2dist, vector<int>(n)), dep(n), sub(n, 1) {
		dfs(0, -1);
	}
 
	/** @return d-th ancestor of u */
	int jump(int u, int d) {
		for (int i = log2dist - 1; i >= 0; i--) {
			if ((d >> i) & 1) { u = lift[i][u]; }
		}
		return u;
	}
 
	/** @return LCA of nodes u and v */
	int lca(int u, int v) {
		if (dep[u] < dep[v]) { swap(u, v); }
		u = jump(u, dep[u] - dep[v]);
		if (u == v) { return u; }
		for (int i = log2dist - 1; i >= 0; i--) {
			if (lift[i][u] != lift[i][v]) {
				u = lift[i][u];
				v = lift[i][v];
			}
		}
		return lift[0][u];
	}
	/** @return # of suitable roots for this query */
	int query(int a, int b, int c) {
		int anc = lca(a, b);
		if (c == anc) {
			return n - sub[jump(a, dep[a] - dep[anc] - 1)] -
			       sub[jump(b, dep[b] - dep[anc] - 1)];
		}
 
		if (dep[c] < dep[anc]) { return 0; }
		if (dep[c] > max(dep[a], dep[b])) { return 0; }
 
		int low = -1;  // child of c that's on the path
		for (int x : {a, b}) {
			if (jump(x, dep[x] - dep[c]) == c) { low = jump(x, dep[x] - dep[c] - 1); }
		}
 
		if (low == -1) { return 0; }
		return sub[c] - sub[low];
	}
};
 
int main() {
	int n, q;
	cin >> n >> q;
	vector<vector<int>> adj(n);
	for (int i = 0; i < n - 1; i++) {
		int u, v;
		cin >> u >> v;
		adj[--u].push_back(--v);
		adj[v].push_back(u);
	}
 
	Tree tree(adj);
	for (int i = 0; i < q; i++) {
		int a, b, c;
		cin >> a >> b >> c;
		cout << tree.query(--a, --b, --c) << "\n";
	}
}