Kahn’s algorithm for Topological Sorting

sort.go

package sort

import (
	"errors"
)

func TopologicalSort(g *Graph) ([]int, error) {
	inDegree := make([]int, g.numOfVertices)

	for _, l := range g.graph {
		for _, j := range l {
			inDegree[j]++
		}
	}

	q := newQueue()

	for i := 0; i < g.numOfVertices; i++ {
		if inDegree[i] == 0 {
			q.push(i)
		}
	}

	count := 0

	topOrder := make([]int, g.numOfVertices)

	for q.size() > 0 {
		u, _ := q.poll()

		topOrder[count] = u

		for _, i := range g.graph[u] {
			inDegree[i]--
			if inDegree[i] == 0 {
				q.push(i)
			}
		}

		count++
	}

	if count != g.numOfVertices {
		return nil, errors.New("There exists a cycle in the graph")
	}

	return topOrder, nil
}

type Graph struct {
	graph         map[int][]int
	numOfVertices int
}

func NewGraph(numOfVertices int) *Graph {
	return &Graph{graph: make(map[int][]int), numOfVertices: numOfVertices}
}

func (g *Graph) AddEdge(u int, v int) {
	g.graph[u] = append(g.graph[u], v)
}

type queue struct {
	queue []int
	pos   int
}

func newQueue() *queue {
	return &queue{queue: make([]int, 10), pos: 0}
}

func (q *queue) push(v int) {
	q.queue[q.pos] = v
	q.pos++
}

func (q *queue) poll() (int, error) {
	if q.pos == 0 {
		return 0, errors.New("cannot poll empty queue")
	}
	var res int
	res, q.queue = q.queue[0], q.queue[1:]
	q.pos--
	return res, nil
}

func (q *queue) size() int {
	return q.pos
}

sort_test.go

package sort

import (
	"testing"
)

func TestTopologicalSort(t *testing.T) {

	g := NewGraph(6)

	g.AddEdge(5, 2)
	g.AddEdge(5, 0)
	g.AddEdge(4, 0)
	g.AddEdge(4, 1)
	g.AddEdge(2, 3)
	g.AddEdge(3, 1)

	sorted, err := TopologicalSort(g)

	if err != nil {
		t.Error("Sort failed")
	}

	expected := []int{4, 5, 2, 0, 3, 1}

	for i, v := range expected {
		if sorted[i] != v {
			t.Error("Sort is wrong")
		}
	}
}

func TestPush(t *testing.T) {
	q := newQueue()

	q.push(2)
	q.push(6)
	q.push(-1)

	if q.size() != 3 {
		t.Errorf("Queue has wrong size, got %d want %d", q.size(), 3)
	}
}

func TestPoll(t *testing.T) {
	q := newQueue()

	tables := []struct {
		expected int
	}{
		{2},
		{6},
		{-1},
	}

	for _, table := range tables {
		q.push(table.expected)
	}

	if q.size() != 3 {
		t.Errorf("Queue has wrong size, got %d want %d", q.size(), 3)
	}

	for _, table := range tables {
		res, _ := q.poll()

		if res != table.expected {
			t.Errorf("Wrong result for poll, got %d want %d", res, table.expected)
		}
	}

	_, err := q.poll()

	if err == nil {
		t.Errorf("Error excpeted when polling am empty queue")
	}
}