Simple Shortest Paths in Java

dijkstra.java

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.NavigableSet;
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.IntStream;
import static java.util.Comparator.comparing;

public class Dijkstra {
  static class Edge {
    final double length;
    final int to;
    Edge(double l, int t) { length = l; to = t; }
  }

  static class SolutionRow {
    double distance;
    int previous = -1;
    SolutionRow(double d) { distance = d; }

    @Override
    public String toString() { return "[" + distance + ": " + previous + "]"; }
  }

  private final List<Set<Edge>> graph = new ArrayList<>();
  private final Set<Integer> visited = new HashSet<>();
  private final List<SolutionRow> solution = new ArrayList<>();
  private final NavigableSet<Integer> unvisited = new TreeSet<>(comparing(i -> solution.get(i).distance));

  private void initialize() {
    unvisited.clear();
    visited.clear();
    solution.clear();
    IntStream.range(0, graph.size()).forEach(i -> {
      solution.add(new SolutionRow(i == 0 ? 0 : Double.POSITIVE_INFINITY));
      unvisited.add(i);
    });
  }

  Dijkstra solve() {
    initialize();
    while (!unvisited.isEmpty()) {
      int currentVertex = unvisited.pollFirst();
      for (Edge edge : graph.get(currentVertex)) {
        int neighbor = edge.to;
        if (!visited.contains(neighbor)) {
          double neighborDistance = solution.get(currentVertex).distance + edge.length;
          SolutionRow neighborSolutionRow = solution.get(neighbor);
          if (neighborDistance < neighborSolutionRow.distance) {
            unvisited.remove(neighbor); // Decrease key simulation.
            neighborSolutionRow.distance = neighborDistance;
            neighborSolutionRow.previous = currentVertex;
            unvisited.add(neighbor);
          }
        }
      }
      visited.add(currentVertex);
    }
    return this;
  }

  List<SolutionRow> getSolution() {
    return solution;
  }

  Dijkstra addUndirected(double len, int a, int b) {
    addDirected(len, a, b);
    return addDirected(len, b, a);
  }

  Dijkstra addDirected(double len, int from, int to) {
    IntStream.rangeClosed(graph.size(), from).forEach(i -> graph.add(new HashSet<>()));
    graph.get(from).add(new Edge(len, to));
    return this;
  }

  public static void main(String[] args) {
    new Dijkstra()
        .addUndirected(6, 0, 1)
        .addUndirected(1, 0, 3)
        .addUndirected(2, 1, 3)
        .addUndirected(2, 1, 4)
        .addUndirected(5, 1, 2)
        .addUndirected(5, 2, 4)
        .addUndirected(1, 3, 4)
        .solve()
        .getSolution().forEach(row -> System.out.println(row));
  }
}