diff --git a/.gitignore b/.gitignore index 3b735ec..82c4ca1 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,5 @@ # Go workspace file go.work + +.idea \ No newline at end of file diff --git a/README.md b/README.md index eb4c813..d39fb99 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,47 @@ # go-prioritize -Priority Queue for Go +Priority Queue for Go with Generic support + +## Installation + +```bash +go get github.com/jinwoo1225/go-prioritize +``` + +## Usage + +```go +package main + +import ( + "fmt" + "github.com/jinwoo1225/go-prioritize" +) + +func main() { + // Create a new priority queue + pq := prioritize.NewPriorityQueue[string](nil) + + // Add some items + pq.Push(prioritize.NewItem[string]("foo", 1)) + pq.Push(prioritize.NewItem[string]("bar", 2)) + pq.Push(prioritize.NewItem[string]("baz", 3)) + + // Pop the highest priority item + item := pq.Pop() + fmt.Println(item) // "baz" + + // Pop the next highest priority item + item = pq.Pop() + fmt.Println(item) // "bar" + + // Pop the next highest priority item + item = pq.Pop() + fmt.Println(item) // "foo" + + // Pop the next highest priority item + item = pq.Pop() + fmt.Println(item) // nil (queue is empty) +} + +``` + diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..b353167 --- /dev/null +++ b/go.mod @@ -0,0 +1,11 @@ +module github.com/jinwoo1225/go-prioritize + +go 1.21.0 + +require github.com/stretchr/testify v1.8.4 + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..fa4b6e6 --- /dev/null +++ b/go.sum @@ -0,0 +1,10 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal_heap.go b/internal_heap.go new file mode 100644 index 0000000..cd405c4 --- /dev/null +++ b/internal_heap.go @@ -0,0 +1,41 @@ +package prioritize + +import "container/heap" + +type internalHeap[T any] struct { + items []Item[T] +} + +var _ heap.Interface = &internalHeap[any]{} + +func (h *internalHeap[T]) Len() int { + return len(h.items) +} + +func (h *internalHeap[T]) Less(i, j int) bool { + return h.items[i].Priority > h.items[j].Priority +} + +func (h *internalHeap[T]) Swap(i, j int) { + h.items[i], h.items[j] = h.items[j], h.items[i] + h.items[i].index = i + h.items[j].index = j +} + +func (h *internalHeap[T]) Push(x interface{}) { + item := x.(Item[T]) + item.index = len(h.items) + h.items = append(h.items, item) +} + +func (h *internalHeap[T]) Pop() interface{} { + old := h.items + n := len(old) + item := old[0] + h.items = old[1:n] + return item +} + +func (h *internalHeap[T]) Fix(i int) { + heap.Fix(h, i) +} diff --git a/item.go b/item.go new file mode 100644 index 0000000..2655b84 --- /dev/null +++ b/item.go @@ -0,0 +1,20 @@ +package prioritize + +// Item for generic support +type Item[T any] struct { + // The Value of the item; arbitrary. + Value T + // The Priority of the item in the queue. + Priority int64 + // The index of the item in the heap. + // This is needed by update and is maintained by the heap.Interface methods. + index int +} + +func NewItem[T any](value T, priority int64) Item[T] { + return Item[T]{ + Value: value, + Priority: priority, + index: -1, + } +} diff --git a/priority_queue.go b/priority_queue.go new file mode 100644 index 0000000..6b58e43 --- /dev/null +++ b/priority_queue.go @@ -0,0 +1,85 @@ +package prioritize + +import ( + "container/heap" + "sync" +) + +type PriorityQueuer[T any] interface { + Pop() Item[T] + Push(x Item[T]) + Update(item Item[T], value T, priority int64) + Peek() Item[T] + Range() []Item[T] +} + +type PriorityQueue[T any] struct { + mu sync.RWMutex + internalHeap *internalHeap[T] +} + +func NewPriorityQueue[T any](items []Item[T]) *PriorityQueue[T] { + for i := range items { + items[i].index = i + } + + iHeap := &internalHeap[T]{ + items: items, + } + + heap.Init(iHeap) + + return &PriorityQueue[T]{ + mu: sync.RWMutex{}, + internalHeap: iHeap, + } +} + +func (h *PriorityQueue[T]) Push(x Item[T]) { + h.mu.Lock() + defer h.mu.Unlock() + + h.internalHeap.Push(x) + h.internalHeap.Fix(h.internalHeap.Len() - 1) +} + +func (h *PriorityQueue[T]) Pop() Item[T] { + h.mu.Lock() + defer h.mu.Unlock() + + return h.internalHeap.Pop().(Item[T]) +} + +func (h *PriorityQueue[T]) Peek() Item[T] { + h.mu.RLock() + defer h.mu.RUnlock() + + if h.internalHeap.Len() == 0 { + return Item[T]{} + } + + return h.internalHeap.items[0] +} + +func (h *PriorityQueue[T]) Update(item Item[T], value T, priority int64) { + h.mu.Lock() + defer h.mu.Unlock() + + item.Value = value + item.Priority = priority + + h.internalHeap.items[item.index] = item + h.internalHeap.Fix(item.index) +} + +func (h *PriorityQueue[T]) Range() []Item[T] { + h.mu.RLock() + defer h.mu.RUnlock() + + var items []Item[T] + for _, item := range h.internalHeap.items { + items = append(items, item) + } + + return items +} diff --git a/priority_queue_test.go b/priority_queue_test.go new file mode 100644 index 0000000..e83b910 --- /dev/null +++ b/priority_queue_test.go @@ -0,0 +1,188 @@ +package prioritize + +import ( + "github.com/stretchr/testify/assert" + "sync" + "testing" +) + +func TestNewHeapedQueue(t *testing.T) { + t.Run("success", func(t *testing.T) { + items := []Item[string]{ + NewItem[string]("hello world", 4), + } + + got := NewPriorityQueue[string](items) + + want := &PriorityQueue[string]{ + mu: sync.RWMutex{}, + internalHeap: &internalHeap[string]{items: items}, + } + + assert.Equal(t, want, got) + }) +} + +func TestPriorityQueue_Push(t *testing.T) { + t.Run("success", func(t *testing.T) { + items := []Item[string]{ + NewItem[string]("hello world", 4), + NewItem[string]("hello world 2", 3), + NewItem[string]("hello world 3", 2), + } + + pq := NewPriorityQueue[string](items) + pq.Push(NewItem[string]("hello world 4", 1)) + + got := pq.internalHeap.items + + want := []Item[string]{ + { + Value: "hello world", + Priority: 4, + index: 0, + }, + { + Value: "hello world 2", + Priority: 3, + index: 1, + }, + { + Value: "hello world 3", + Priority: 2, + index: 2, + }, + { + Value: "hello world 4", + Priority: 1, + index: 3, + }, + } + + assert.Equal(t, want, got) + }) +} + +func TestPriorityQueue_Pop(t *testing.T) { + t.Run("success", func(t *testing.T) { + items := []Item[string]{ + NewItem[string]("hello world", 4), + NewItem[string]("hello world 2", 3), + NewItem[string]("hello world 3", 2), + NewItem[string]("hello world 4", 1), + } + + pq := NewPriorityQueue[string](items) + got := pq.Pop() + + want := Item[string]{ + Value: "hello world", + Priority: 4, + index: 0, + } + + assert.Equal(t, want, got) + }) +} + +func TestPriorityQueue_Peek(t *testing.T) { + t.Run("success", func(t *testing.T) { + items := []Item[string]{ + NewItem[string]("hello world", 4), + NewItem[string]("hello world 2", 3), + NewItem[string]("hello world 3", 2), + NewItem[string]("hello world 4", 1), + } + + pq := NewPriorityQueue[string](items) + got := pq.Peek() + + want := Item[string]{ + Value: "hello world", + Priority: 4, + index: 0, + } + + assert.Equal(t, want, got) + }) +} + +func TestPriorityQueue_Update(t *testing.T) { + t.Run("success", func(t *testing.T) { + items := []Item[string]{ + NewItem[string]("hello world", 4), + NewItem[string]("hello world 2", 3), + NewItem[string]("hello world 3", 2), + NewItem[string]("hello world 4", 1), + } + + pq := NewPriorityQueue[string](items) + pq.Update(pq.internalHeap.items[0], "hello world updated", 5) + + got := pq.internalHeap.items + + want := []Item[string]{ + { + Value: "hello world updated", + Priority: 5, + index: 0, + }, + { + Value: "hello world 2", + Priority: 3, + index: 1, + }, + { + Value: "hello world 3", + Priority: 2, + index: 2, + }, + { + Value: "hello world 4", + Priority: 1, + index: 3, + }, + } + + assert.Equal(t, want, got) + }) +} + +func TestPriorityQueue_Range(t *testing.T) { + t.Run("success", func(t *testing.T) { + items := []Item[string]{ + NewItem[string]("hello world", 4), + NewItem[string]("hello world 2", 3), + NewItem[string]("hello world 3", 2), + NewItem[string]("hello world 4", 1), + } + + pq := NewPriorityQueue[string](items) + got := pq.Range() + + want := []Item[string]{ + { + Value: "hello world", + Priority: 4, + index: 0, + }, + { + Value: "hello world 2", + Priority: 3, + index: 1, + }, + { + Value: "hello world 3", + Priority: 2, + index: 2, + }, + { + Value: "hello world 4", + Priority: 1, + index: 3, + }, + } + + assert.Equal(t, want, got) + }) +}