diff --git a/README.md b/README.md
index 0e61cf7..74c189b 100644
--- a/README.md
+++ b/README.md
@@ -30,6 +30,7 @@ You can add the flag `-benchmem` if you want to see the memory allocations:
 go test ./... -bench=. -run=xxx -benchmem
 ```
 
+
 ## Notes:
 
 * Maps are one of the common ways to store graphs in Go. The `TopologicalSort` function input supports `map[interface{}][]interface{}` maps.
diff --git a/cmd/example-2/example_2.go b/cmd/example-2/example_2.go
index 0c6dd8b..5c648b6 100644
--- a/cmd/example-2/example_2.go
+++ b/cmd/example-2/example_2.go
@@ -37,13 +37,13 @@ func main() {
 		3: []int{4, 5},
 	}
 
-	output, err = gorder.TopologicalSort(serialize(digraph), "kahn")
+	output, err = gorder.TopologicalSort(serialize(digraph), gorder.KAHN)
 	if err != nil {
 		log.Fatal(err)
 	}
 	fmt.Printf("Solution (Kahn): %d\n", deserialize(output))
 
-	output, err = gorder.TopologicalSort(serialize(digraph), "dfsbased")
+	output, err = gorder.TopologicalSort(serialize(digraph), gorder.DFS)
 	if err != nil {
 		log.Fatal(err)
 	}
diff --git a/cmd/example/example.go b/cmd/example/example.go
index e3902c8..64c5bd5 100644
--- a/cmd/example/example.go
+++ b/cmd/example/example.go
@@ -23,13 +23,13 @@ func main() {
 		3: []interface{}{"4", "5"},
 	}
 
-	o, err = gorder.TopologicalSort(digraph, "kahn")
+	o, err = gorder.TopologicalSort(digraph, gorder.KAHN)
 	if err != nil {
 		log.Fatal(err)
 	}
 	fmt.Printf("Solution (Kahn): %v\n", o)
 
-	o, err = gorder.TopologicalSort(digraph, "dfsbased")
+	o, err = gorder.TopologicalSort(digraph, gorder.DFS)
 	if err != nil {
 		log.Fatal(err)
 	}
diff --git a/cmd/measurements/measurements.go b/cmd/measurements/measurements.go
index c02466f..8491d39 100644
--- a/cmd/measurements/measurements.go
+++ b/cmd/measurements/measurements.go
@@ -15,7 +15,7 @@ func main() {
 	log.Printf("DAG generation time: %v", time.Since(start))
 
 	start = time.Now()
-	s, err := gorder.TopologicalSort(graph, "kahn")
+	s, err := gorder.TopologicalSort(graph, gorder.KAHN)
 	log.Printf("Kahn resolve time: %v", time.Since(start))
 	if err != nil {
 		log.Fatal(err)
@@ -23,7 +23,7 @@ func main() {
 	spew.Dump(s)
 
 	start = time.Now()
-	s, err = gorder.TopologicalSort(graph, "dfsbased")
+	s, err = gorder.TopologicalSort(graph, gorder.DFS)
 	log.Printf("DFS-based resolve time: %v", time.Since(start))
 	if err != nil {
 		log.Fatal(err)
diff --git a/go.mod b/go.mod
index 0c34830..b531060 100644
--- a/go.mod
+++ b/go.mod
@@ -1,5 +1,5 @@
 module github.com/amwolff/gorder
 
-go 1.16
+go 1.20
 
 require github.com/davecgh/go-spew v1.1.1
diff --git a/gorder.go b/gorder.go
index 4bad6e4..6407155 100644
--- a/gorder.go
+++ b/gorder.go
@@ -2,24 +2,30 @@ package gorder
 
 import (
 	"errors"
-	"regexp"
 )
 
-func TopologicalSort(digraph map[interface{}][]interface{}, algorithm string) (solution []interface{}, err error) {
-	kahnRgxp, err := regexp.Compile(`[Kk]ahn\z`)
-	if err != nil {
-		return nil, err
-	}
-	dfsBasedRgxp, err := regexp.Compile(`[Dd][Ff][Ss]-?[Bb]ased\z`)
-	if err != nil {
-		return nil, err
-	}
+type algo int
+
+const (
+	DFS algo = iota
+	KAHN
+)
+
+func TopologicalSort[T comparable, V []T](digraph map[T]V, algorithm algo) (solution V, err error) {
+	// kahnRgxp, err := regexp.Compile(`[Kk]ahn\z`)
+	// if err != nil {
+	// 	return nil, err
+	// }
+	// dfsBasedRgxp, err := regexp.Compile(`[Dd][Ff][Ss]-?[Bb]ased\z`)
+	// if err != nil {
+	// 	return nil, err
+	// }
 
-	if kahnRgxp.MatchString(algorithm) {
+	if algorithm == KAHN {
 		if solution, err = kahn(digraph); err != nil {
 			return nil, err
 		}
-	} else if dfsBasedRgxp.MatchString(algorithm) {
+	} else if algorithm == DFS {
 		if solution, err = dfsBased(digraph); err != nil {
 			return nil, err
 		}
@@ -29,24 +35,24 @@ func TopologicalSort(digraph map[interface{}][]interface{}, algorithm string) (s
 	return solution, nil
 }
 
-func kahn(digraph map[interface{}][]interface{}) ([]interface{}, error) {
-	indegrees := make(map[interface{}]int)
-	for u := range digraph {
-		if digraph[u] != nil {
-			for _, v := range digraph[u] {
-				indegrees[v]++
-			}
+func kahn[T comparable, V []T](digraph map[T]V) (V, error) {
+	indegrees := make(map[T]int)
+
+	// loop through all diagraph and add increase indegrees of values
+	for _, iter := range digraph {
+		for _, v := range iter {
+			indegrees[v]++
 		}
 	}
 
-	var queue []interface{}
+	var queue V
 	for u := range digraph {
 		if _, ok := indegrees[u]; !ok {
 			queue = append(queue, u)
 		}
 	}
 
-	var order []interface{}
+	var order V
 	for len(queue) > 0 {
 		u := queue[len(queue)-1]
 		queue = queue[:(len(queue) - 1)]
@@ -67,16 +73,16 @@ func kahn(digraph map[interface{}][]interface{}) ([]interface{}, error) {
 	return order, nil
 }
 
-func dfsBased(digraph map[interface{}][]interface{}) ([]interface{}, error) {
+func dfsBased[T comparable, V []T](digraph map[T]V) (V, error) {
 	var (
 		acyclic       = true
-		order         []interface{}
-		permanentMark = make(map[interface{}]bool)
-		temporaryMark = make(map[interface{}]bool)
-		visit         func(interface{})
+		order         V
+		permanentMark = make(map[T]bool)
+		temporaryMark = make(map[T]bool)
+		visit         func(T)
 	)
 
-	visit = func(u interface{}) {
+	visit = func(u T) {
 		if temporaryMark[u] {
 			acyclic = false
 		} else if !(temporaryMark[u] || permanentMark[u]) {
@@ -89,7 +95,7 @@ func dfsBased(digraph map[interface{}][]interface{}) ([]interface{}, error) {
 			}
 			delete(temporaryMark, u)
 			permanentMark[u] = true
-			order = append([]interface{}{u}, order...)
+			order = append(V{u}, order...)
 		}
 	}
 
diff --git a/gorder_test.go b/gorder_test.go
index b2a8f3e..35ec941 100644
--- a/gorder_test.go
+++ b/gorder_test.go
@@ -7,20 +7,21 @@ import (
 )
 
 func TestTopologicalSort(t *testing.T) {
-	digraph := map[interface{}][]interface{}{
-		1: []interface{}{2, 4},
-		2: []interface{}{3, 5},
-		3: []interface{}{4, 5},
+	digraph := map[int][]int{
+		1: {2, 4},
+		2: {3, 5},
+		3: {4, 5},
 	}
 
 	want := []int{1, 2, 3, 5, 4}
 
-	_, err := TopologicalSort(digraph, "ultrasuperfast")
+
+	_, err := TopologicalSort(digraph, 3)
 	if err == nil {
 		t.Fatal("TopologicalSort: should have returned an error")
 	}
 
-	o, err := TopologicalSort(digraph, "kahn")
+	o, err := TopologicalSort(digraph, KAHN)
 	if err != nil {
 		t.Fatalf("Kahn: %v", err)
 	}
@@ -30,7 +31,7 @@ func TestTopologicalSort(t *testing.T) {
 		}
 	}
 
-	o, err = TopologicalSort(digraph, "dfsbased")
+	o, err = TopologicalSort(digraph, DFS)
 	if err != nil {
 		t.Fatalf("DFS-based: %v", err)
 	}
@@ -40,17 +41,17 @@ func TestTopologicalSort(t *testing.T) {
 		}
 	}
 
-	graphWithCycles := map[interface{}][]interface{}{
-		1: []interface{}{2, 4},
-		2: []interface{}{3, 5},
-		3: []interface{}{4, 5},
-		4: []interface{}{2},
+	graphWithCycles := map[int][]int{
+		1: {2, 4},
+		2: {3, 5},
+		3: {4, 5},
+		4: {2},
 	}
-	_, err = TopologicalSort(graphWithCycles, "kahn")
+	_, err = TopologicalSort(graphWithCycles, KAHN)
 	if err == nil {
 		t.Fatal("Kahn: should have returned an error")
 	}
-	_, err = TopologicalSort(graphWithCycles, "dfsbased")
+	_, err = TopologicalSort(graphWithCycles, DFS)
 	if err == nil {
 		t.Fatal("DFS-based: should have returned an error")
 	}
@@ -60,7 +61,7 @@ func BenchmarkTopologicalSort(b *testing.B) {
 	digraph := dagenerator.Generate(10, 50, 30, 50, 30)
 	b.Run("kahn", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
-			_, err := TopologicalSort(digraph, "kahn")
+			_, err := TopologicalSort(digraph, KAHN)
 			if err != nil {
 				b.Errorf("Kahn: %v", err)
 			}
@@ -68,7 +69,7 @@ func BenchmarkTopologicalSort(b *testing.B) {
 	})
 	b.Run("dfs based", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
-			_, err := TopologicalSort(digraph, "dfsbased")
+			_, err := TopologicalSort(digraph, DFS)
 			if err != nil {
 				b.Errorf("DFS-based: %v", err)
 			}