Tree Diameter

Solución en video

Por Abhiraj Mallangi

Nota: la solución en video puede no coincidir con las demás soluciones.

Solución 1 - DFS dos veces

Explicación

Ejecutamos DFS desde un nodo arbitrario y calculamos la distancia a cada nodo. Como el diámetro es el camino más largo, uno de sus extremos debe ser el más lejano del punto de partida. Por lo tanto, el nodo con distancia máxima es uno de los extremos del diámetro. Podemos ejecutar un segundo DFS desde este extremo, pues la distancia máxima hallada será el diámetro.

Este enfoque es el Approach 2 en CPH 14.2.

Implementación

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

C++JavaPython


#include <bits/stdc++.h>
using namespace std;
 
vector<int> dist;
vector<vector<int>> adj;
 
void dfs(int curr, int parent) {
	if (parent != -1) { dist[curr] = dist[parent] + 1; }
	for (int next : adj[curr]) {
		if (next != parent) { dfs(next, curr); }
	}
}
 
int main() {
	int n;
	cin >> n;
	dist.resize(n);
	adj.resize(n);
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		cin >> a >> b;
		adj[--a].push_back(--b);
		adj[b].push_back(a);
	}
 
	// first DFS to find furthest node from node 0
	dfs(0, -1);
 
	int start = -1, maxDist = -1;
	for (int i = 0; i < n; i++) {
		if (dist[i] > maxDist) {
			maxDist = dist[i];
			start = i;
		}
	}
 
	// second DFS to find other endpoint of diameter
	fill(dist.begin(), dist.end(), 0);
	dfs(start, -1);
 
	maxDist = -1;
	for (int i = 0; i < n; i++) { maxDist = max(maxDist, dist[i]); }
	cout << maxDist << "\n";
}


import java.io.*;
import java.util.*;
 
public class TreeDiameter {
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
 
		int n = Integer.parseInt(br.readLine());
		int[] dist = new int[n];
		ArrayList<Integer>[] adj = new ArrayList[n];
		for (int i = 0; i < n; i++) { adj[i] = new ArrayList<>(); }
		for (int i = 0; i < n - 1; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			int a = Integer.parseInt(st.nextToken()) - 1;
			int b = Integer.parseInt(st.nextToken()) - 1;
			adj[a].add(b);
			adj[b].add(a);
		}
 
		// First DFS from node 0
		Arrays.fill(dist, -1);
		dist[0] = 0;
		Deque<Integer> stack = new ArrayDeque<>();
		stack.push(0);
		while (!stack.isEmpty()) {
			int curr = stack.pop();
			for (int next : adj[curr]) {
				if (dist[next] == -1) {
					dist[next] = dist[curr] + 1;
					stack.push(next);
				}
			}
		}
 
		int start = 0;
		for (int i = 0; i < n; i++) {
			if (dist[i] > dist[start]) start = i;
		}
 
		// Second DFS from farthest node
		Arrays.fill(dist, -1);
		dist[start] = 0;
		stack.push(start);
		while (!stack.isEmpty()) {
			int curr = stack.pop();
			for (int next : adj[curr]) {
				if (dist[next] == -1) {
					dist[next] = dist[curr] + 1;
					stack.push(next);
				}
			}
		}
 
		int maxDist = -1;
		for (int i = 0; i < n; i++) { maxDist = Math.max(maxDist, dist[i]); }
 
		System.out.println(maxDist);
	}
}


import sys
 
input = sys.stdin.readline  # redefine input for performance reasons
sys.setrecursionlimit(10**7)
 
# iterative DFS necessary in order to AC; recursive TLEs
def iterative_dfs(start):
	dist = [-1] * n
	stack = [start]
	dist[start] = 0
	while stack:
		curr = stack.pop()
		for nxt in adj[curr]:
			if dist[nxt] == -1:
				dist[nxt] = dist[curr] + 1
				stack.append(nxt)
	return dist
 
 
n = int(input())
adj = [[] for _ in range(n)]
for _ in range(n - 1):
	a, b = map(int, input().split())
	a -= 1
	b -= 1
	adj[a].append(b)
	adj[b].append(a)
 
# dfs from node 0 to find endpoint of diameter
dist = iterative_dfs(0)
start = max(range(n), key=lambda i: dist[i])
 
# dfs from endpoint of diameter to find other endpoint
dist = iterative_dfs(start)
print(max(dist))

Solución 2 - DP en árboles

Explicación

Primero enraizamos el árbol de forma arbitraria para establecer un orden de los nodos. Con DFS, recorremos hacia abajo desde la raíz y computamos las alturas de los dos subárboles más altos de cada nodo. Sumándolas de forma bottom-up, hallamos la longitud del camino más largo que pasa por ese punto. El diámetro es el máximo de los caminos más largos que hemos hallado en todos los nodos.

Este enfoque de DP es el Approach 1 en CPH.

Implementación

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

C++JavaPython


#include <bits/stdc++.h>
using namespace std;
 
int diameter = 0;
vector<vector<int>> adj;
 
int dfs(int curr, int parent) {
	int max1 = 0, max2 = 0;
	for (int next : adj[curr]) {
		if (next != parent) {
			int h = dfs(next, curr) + 1;
			if (h > max1) {
				max2 = max1;
				max1 = h;
			} else if (h > max2) {
				max2 = h;
			}
		}
	}
 
	// diameter will equal the max of the sum of the heights of the two tallest subtrees
	diameter = max(diameter, max1 + max2);
	return max1;  // max1 is the height of the current subtree
}
 
int main() {
	int n;
	cin >> n;
	adj.resize(n);
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		cin >> a >> b;
		adj[--a].push_back(--b);
		adj[b].push_back(a);
	}
 
	dfs(0, -1);
 
	cout << diameter << endl;
}


import java.io.*;
import java.util.*;
 
public class TreeDiameter {
	private static ArrayList<Integer>[] adj;
	private static int[] height;
	private static int diameter = 0;
 
	static class State {
		public int node, parent, phase;
		State(int node, int parent, int phase) {
			this.node = node;
			this.parent = parent;
			this.phase = phase;
		}
	}
 
	// iterative DFS necessary in order to AC; recursive TLEs on final case
	static void calculateDiameter(int start) {
		Deque<State> stack = new ArrayDeque<>();
		// start from node 0, no parent, pre-processing phase
		stack.push(new State(start, -1, 0));
 
		while (!stack.isEmpty()) {
			State curr = stack.pop();
			int u = curr.node;
			int p = curr.parent;
			int phase = curr.phase;
 
			if (phase == 0) {
				// convert current state to post-processing
				stack.push(new State(u, p, 1));
 
				for (int v : adj[u]) {
					if (v != p) { stack.push(new State(v, u, 0)); }
				}
			} else {
				// post-processing: calculate heights and update diameter
				int max1 = 0, max2 = 0;
				for (int v : adj[u]) {
					if (v == p) continue;
					int h = height[v] + 1;
					if (h > max1) {
						max2 = max1;
						max1 = h;
					} else if (h > max2) {
						max2 = h;
					}
				}
				height[u] = max1;
				// diameter will equal the max of the sum of the heights of the two
				// tallest subtrees
				diameter = Math.max(diameter, max1 + max2);
			}
		}
	}
 
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int n = Integer.parseInt(br.readLine());
 
		adj = new ArrayList[n];
		for (int i = 0; i < n; i++) { adj[i] = new ArrayList<>(); }
		for (int i = 0; i < n - 1; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			int a = Integer.parseInt(st.nextToken()) - 1;
			int b = Integer.parseInt(st.nextToken()) - 1;
			adj[a].add(b);
			adj[b].add(a);
		}
 
		height = new int[n];
		diameter = 0;
		calculateDiameter(0);
 
		System.out.println(diameter);
	}
}


import sys
 
input = sys.stdin.readline  # redefine input for performance reasons
sys.setrecursionlimit(10**7)
 
n = int(input())
adj = [[] for _ in range(n)]
for _ in range(n - 1):
	a, b = map(int, input().split())
	a -= 1
	b -= 1
	adj[a].append(b)
	adj[b].append(a)
 
height = [0] * n
diameter = 0
 
# iterative DFS necessary in order to AC; recursive TLEs
stack = [(0, -1, 0)]  # node, parent, state (0=pre, 1=post)
 
while stack:
	u, p, state = stack.pop()
	if state == 0:
		stack.append((u, p, 1))  # convert current state to post-processing
		for v in adj[u]:
			if v != p:
				stack.append((v, u, 0))  # pre-process children
	else:
		# post-processing: after processing all children, recalculate max heights of subtrees
		max1 = max2 = 0
		for v in adj[u]:
			if v == p:
				continue
			h = height[v] + 1
			if h > max1:
				max2 = max1
				max1 = h
			elif h > max2:
				max2 = h
		height[u] = max1
		# diameter will equal the max of the sum of the heights of the two tallest subtrees
		diameter = max(diameter, max1 + max2)
 
print(diameter)