diff --git a/cmd/internal/graph/common.go b/cmd/internal/graph/common.go
index 2783b83..39dca4e 100644
--- a/cmd/internal/graph/common.go
+++ b/cmd/internal/graph/common.go
@@ -40,8 +40,7 @@ type WeightedMatrix interface {
 
 func toAdjMatrix(g WeightedGraph, adj WeightedMatrix) *mat.Dense {
 	copyEdges(g, adj)
-	matrix := mat.DenseCopyOf(adj.Matrix())
-	addSelfEdges(g, matrix)
+	matrix := addSelfEdges(g, adj)
 
 	return matrix
 }
@@ -57,14 +56,19 @@ func copyEdges(g WeightedGraph, adj WeightedMatrix) {
 	}
 }
 
-func addSelfEdges(g WeightedGraph, matrix mat.Mutable) {
-	nodes := g.Nodes()
+func addSelfEdges(g WeightedGraph, adj WeightedMatrix) *mat.Dense {
+	matrix := mat.DenseCopyOf(adj.Matrix())
+
+	nodes := adj.Nodes()
 	for i := 0; nodes.Next(); i++ {
-		u := nodes.Node()
+		id := nodes.Node().ID()
+		u := g.Node(id)
 
 		e := g.WeightedEdge(u.ID(), u.ID())
 		if e != nil {
 			matrix.Set(i, i, e.Weight())
 		}
 	}
+
+	return matrix
 }