noic-solucao-intermediario-s57.cpp

#include <bits/stdc++.h>
#define F first
#define S second

using namespace std;

typedef pair<int, int> pii;
const int maxn = 10010, maxl = 20; // declaro os limites

// declaro as variáveis que vou utilizar
int n, q;
int ancestral[maxn][maxl];
int dist[maxn], nivel[maxn];

vector<pii> graph[maxn];

void dfs(int u, int pai) 
{
	if(pai != -1) nivel[u] = nivel[pai] + 1; // calculo o meu nível

	ancestral[u][0] = p; // coloco meu pai na tabela

	for(auto v: graph[u]) // percorro meus filhos
	{
		if(v.S == pai) continue; // se meu filho é igual o meu pai, continuo
								 // se não

		dist[v.S] = dist[u] + v.F; // calculo minha distância á raiz
		
		dfs(v.S, u); // chamo a dfs
	}
}

void tabela() // monto a tabela dos ancestrais
{
	memset(ancestral, -1, sizeof ancestral);
	memset(dist, 0, sizeof dist);
	memset(nivel, 0, sizeof nivel);

	dfs(1, -1);

	for(int j = 1 ; j <= maxl ; ++j)
	{
		for(int i = 1 ; i <= n ; ++i)
		{
			if(ancestral[i][j - 1] != -1)
			{
				ancestral[i][j] = ancestral[ancestral[i][j - 1]][j - 1];
			}
		}
	}
}

int LCA(int u, int v) // procuro o LCA entre u e v
{
	if(nivel[u] < nivel[v]) swap(u, v);

	for(int i = maxl - 1 ; i >= 0 ; --i)
	{
		if(nivel[u]-(1<<i) >= nivel[v])
		{
			u = ancestral[u][i];
		}
	}

	if(u == v) return u;

	for(int i = maxl - 1 ; i >= 0 ; --i)
	{
		if(ancestral[u][i] != -1 && ancestral[v][i] != ancestral[u][i])
		{
			u = ancestral[u][i];
			v = ancestral[v][i];
		}
	}

	return ancestral[u][0];
}


int find(int u, int k) // procuro o k-ésimo nó no caminho para a raiz
{
	for(int i = maxl - 1 ; i >= 0 ; --i) // passo por todos os niveis da tabela
	{
		if(k - (1<<i) >=0) // se eu ainda não achei o nó
		{	
			u = ancestral[u][i]; // subo 2^i níveis
			k -= (1<<i); // diminuo minha distância ao nó
		}
	}

	return u;
}

int main()
{
	ios::sync_with_stdio(false); cin.tie(0);

	cin >> n;

	for(int i = 1 ; i < n ; ++i) // monto a árvore
	{
		int u, v, c;
		cin >> u >> v >> c;

		graph[u].push_back(pii(c, v));
		graph[v].push_back(pii(c, u));
	}

	tabela(); // monto a tabela dos ancestrais

	string s;

	for(int i = 1 ; i <= q ; ++i)
	{
		int id, u, v;
		
		cin >> u, v;

		if(id == 1) cout << dist[u] + dist[v] - 2*dist[LCA(u, v)] << "\n";
				
		else
		{
			int k;
			cin >> k;
			
			int lca = LCA(u, v);

			int diferenca_niveis = nivel[u] - nivel[lca] + 1; // calculo a diferença entre os níveis do lca e u

			if(k == diferenca_niveis) // significa que o lca de u e v está a k níveis de distancia de u
			{ 
				cout << lca << "\n";
			}
			
			else if(k < diferenca_niveis) // se o lca estiver no caminho de u até o lca
			{			      // procuro quem esta a k - 1 níveis acima de u.

				cout << find(u, k-1) << "\n";
			}

			else // nesse caso, o k-ésimo nó do caminho de u -> v esta depois do lca, ou seja,
			{    // no caminho de v até o lca.

				cout << find(b, diferenca_niveis + nivel[v] - nivel[lca] - k) << "\n";
			}
		}
	}
	return 0;
}