BOJ. Diameter of Tree (1967)

BOJ 1967: Diameter of Tree

The input lists each undirected edge only once. Store it in both endpoints’ adjacency lists, so a traversal can move in either direction; do not add either direction a second time.

In a tree, the unique simple path between any two vertices is their distance. Start at any vertex s and traverse the whole tree, recording each vertex’s distance from s. Let a be a farthest vertex. A farthest endpoint of a diameter can be found this way: take a diameter path and consider where the path from s meets it. If s lies on the path, one of its endpoints is at least as far from s as that meeting point, and a farthest vertex is a diameter endpoint. If s lies off the path, the unique path from s joins the diameter at a vertex; by the tree’s unique paths, one of the two diameter endpoints is farthest (ties are harmless). Thus a is an endpoint of some diameter. A second traversal from a finds the greatest distance from that endpoint, which is the diameter length.

Both traversals use explicit stacks rather than recursion, so a chain cannot overflow the call stack. The distances and answer use long. The n = 1 case has no edges: both traversals stay at vertex 0 and the answer is 0. Each traversal takes O(n) time, and the adjacency lists and traversal arrays use O(n) space.

import java.io.*;
import java.util.*;

public class Main {
  static class Edge {
    final int to;
    final int weight;

    Edge(int to, int weight) {
      this.to = to;
      this.weight = weight;
    }
  }

  static class Farthest {
    final int vertex;
    final long distance;

    Farthest(int vertex, long distance) {
      this.vertex = vertex;
      this.distance = distance;
    }
  }

  static class FastScanner {
    private final InputStream input = System.in;
    private final byte[] buffer = new byte[1 << 16];
    private int length = 0;
    private int pointer = 0;

    private int read() throws IOException {
      if (pointer == length) {
        length = input.read(buffer);
        pointer = 0;
        if (length == -1) return -1;
      }
      return buffer[pointer++];
    }

    int nextInt() throws IOException {
      int c;
      do {
        c = read();
      } while (c <= ' ' && c != -1);

      int value = 0;
      while (c > ' ') {
        value = value * 10 + c - '0';
        c = read();
      }
      return value;
    }
  }

  static Farthest farthest(ArrayList<Edge>[] graph, int start) {
    int n = graph.length;
    int[] parent = new int[n];
    Arrays.fill(parent, -2);
    int[] stack = new int[n];
    long[] distance = new long[n];
    int size = 0;
    stack[size++] = start;
    parent[start] = -1;

    int bestVertex = start;
    long bestDistance = 0;
    while (size > 0) {
      int vertex = stack[--size];
      if (distance[vertex] > bestDistance) {
        bestDistance = distance[vertex];
        bestVertex = vertex;
      }
      for (Edge edge : graph[vertex]) {
        if (edge.to == parent[vertex]) continue;
        parent[edge.to] = vertex;
        distance[edge.to] = distance[vertex] + edge.weight;
        stack[size++] = edge.to;
      }
    }
    return new Farthest(bestVertex, bestDistance);
  }

  public static void main(String[] args) throws IOException {
    FastScanner input = new FastScanner();
    int n = input.nextInt();
    ArrayList<Edge>[] graph = new ArrayList[n];
    for (int i = 0; i < n; i++) graph[i] = new ArrayList<>();

    for (int i = 0; i < n - 1; i++) {
      int a = input.nextInt() - 1;
      int b = input.nextInt() - 1;
      int weight = input.nextInt();
      graph[a].add(new Edge(b, weight));
      graph[b].add(new Edge(a, weight));
    }

    int endpoint = farthest(graph, 0).vertex;
    System.out.println(farthest(graph, endpoint).distance);
  }
}