Kruskal's Minimum Spanning Tree using Java Stream

KruskalMST.java

package greedy;

import java.util.*;
import java.util.stream.IntStream;

/**
 * Created by ganessh on 28/07/16.
 */
class Edge {
    private int source;
    private int destination;
    private int distance;

    public Edge(int source, int destination, int distance) {
        this.source = source;
        this.destination = destination;
        this.distance = distance;
    }

    public int getDistance() {
        return distance;
    }

    public int getSource() {
        return source;
    }

    public int getDestination() {
        return destination;
    }
}

class Vertex {
    private int id;
    private Vertex parent;
    private int distance;

    public Vertex(int id) {
        this.id = id;
    }

    public int getId() {
        return id;
    }

    public void setId(int id) {
        this.id = id;
    }
}

public class KruskalsMST {
    private static Set findVertexSet(List<Set<Integer>> vertexSet, int vertexId) {
        return vertexSet.stream().filter(vs -> vs.contains(vertexId)).findFirst().get();
    }

    public static void main(String[] args) {
        Map<Integer, Vertex> vertices = new HashMap<Integer, Vertex>();
        IntStream.rangeClosed(0, 8).forEach(i -> vertices.put(i, new Vertex(i)));

        List<Edge> edges = new ArrayList<Edge>() {{
            add(new Edge(0, 1, 4));
            add(new Edge(1, 2, 8));
            add(new Edge(2, 3, 7));
            add(new Edge(3, 4, 9));
            add(new Edge(4, 5, 10));
            add(new Edge(5, 6, 2));
            add(new Edge(6, 7, 1));
            add(new Edge(7, 0, 8));
            add(new Edge(1, 7, 11));
            add(new Edge(2, 8, 2));
            add(new Edge(8, 7, 7));
            add(new Edge(8, 6, 6));
            add(new Edge(2, 5, 4));
            add(new Edge(3, 5, 14));
        }};

        List<Set<Integer>> vertexSets = new ArrayList<Set<Integer>>();
        vertices.values().stream().forEach(
                        vertex -> vertexSets.add(new HashSet<Integer>() {{
                                    add(vertex.getId()) ;
                                }})
                );

        Collections.sort(edges, (e1, e2) -> Integer.compare(e1.getDistance(), e2.getDistance()));
        // Print sorted edges
        // edges.stream().forEach(edge -> System.out.print(edge.getDistance() + " "));

        Set<Edge> resultantEdges = new HashSet<Edge>();

        edges.stream().forEach(edge -> {
            Set<Integer> sourceSet = findVertexSet(vertexSets, edge.getSource());
            Set<Integer> destinationSet = findVertexSet(vertexSets, edge.getDestination());
            if (sourceSet != destinationSet) {
                sourceSet.addAll(destinationSet);
                vertexSets.remove(destinationSet);
                resultantEdges.add(edge);
            }
            System.out.println("-------- Procesed a node --------");
            vertexSets.stream().forEach(vs -> {
                System.out.print("VertexSet ");
                vs.stream().forEach(vid -> System.out.print(vid + " "));
                System.out.println();
            });
            resultantEdges.stream().forEach(e -> System.out.println(e.getSource() + " -> " + e.getDestination()));
            System.out.println("---------------------------------");
        });

        System.out.println("Number of edges in result " + resultantEdges.size());
        resultantEdges.stream().forEach(edge -> System.out.println(edge.getSource() + " -> " + edge.getDestination()));

    }
}

Output.md

-------- Procesed a node --------
VertexSet 0 
VertexSet 1 
VertexSet 2 
VertexSet 3 
VertexSet 4 
VertexSet 5 
VertexSet 6 7 
VertexSet 8 
Edges added so far
6 -> 7
---------------------------------
-------- Procesed a node --------
VertexSet 0 
VertexSet 1 
VertexSet 2 
VertexSet 3 
VertexSet 4 
VertexSet 5 6 7 
VertexSet 8 
Edges added so far
6 -> 7
5 -> 6
---------------------------------
-------- Procesed a node --------
VertexSet 0 
VertexSet 1 
VertexSet 2 8 
VertexSet 3 
VertexSet 4 
VertexSet 5 6 7 
Edges added so far
2 -> 8
6 -> 7
5 -> 6
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 
VertexSet 2 8 
VertexSet 3 
VertexSet 4 
VertexSet 5 6 7 
Edges added so far
0 -> 1
2 -> 8
6 -> 7
5 -> 6
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 
VertexSet 2 5 6 7 8 
VertexSet 3 
VertexSet 4 
Edges added so far
2 -> 5
0 -> 1
2 -> 8
6 -> 7
5 -> 6
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 
VertexSet 2 5 6 7 8 
VertexSet 3 
VertexSet 4 
Edges added so far
2 -> 5
0 -> 1
2 -> 8
6 -> 7
5 -> 6
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 
VertexSet 2 3 5 6 7 8 
VertexSet 4 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
2 -> 8
6 -> 7
5 -> 6
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 
VertexSet 2 3 5 6 7 8 
VertexSet 4 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
2 -> 8
6 -> 7
5 -> 6
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 2 3 5 6 7 8 
VertexSet 4 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
2 -> 8
6 -> 7
5 -> 6
1 -> 2
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 2 3 5 6 7 8 
VertexSet 4 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
2 -> 8
6 -> 7
5 -> 6
1 -> 2
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 2 3 4 5 6 7 8 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
3 -> 4
2 -> 8
6 -> 7
5 -> 6
1 -> 2
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 2 3 4 5 6 7 8 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
3 -> 4
2 -> 8
6 -> 7
5 -> 6
1 -> 2
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 2 3 4 5 6 7 8 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
3 -> 4
2 -> 8
6 -> 7
5 -> 6
1 -> 2
---------------------------------
-------- Procesed a node --------
VertexSet 0 1 2 3 4 5 6 7 8 
Edges added so far
2 -> 5
2 -> 3
0 -> 1
3 -> 4
2 -> 8
6 -> 7
5 -> 6
1 -> 2
---------------------------------
Number of edges in result 8
2 -> 5
2 -> 3
0 -> 1
3 -> 4
2 -> 8
6 -> 7
5 -> 6
1 -> 2