From eaf390ee067a732cd5a0cf70555a37d2a63a9860 Mon Sep 17 00:00:00 2001
From: rdleal <rafaeldias.c@gmail.com>
Date: Thu, 7 Dec 2023 23:25:17 -0300
Subject: [PATCH] feat: add Set method. Closes #1

---
 kpq/example_test.go              | 34 +++++++++++++++++++++++++
 kpq/keyed_priority_queue.go      | 27 ++++++++++++++++++--
 kpq/keyed_priority_queue_test.go | 43 ++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 2 deletions(-)

diff --git a/kpq/example_test.go b/kpq/example_test.go
index 32e259f..d0b2304 100644
--- a/kpq/example_test.go
+++ b/kpq/example_test.go
@@ -62,3 +62,37 @@ func Example() {
 	// Key: "key2", Value: 30
 	// Key 'key3' exists: true
 }
+
+func ExampleKeyedPriorityQueue_Set() {
+	cmp := func(a, b int) bool {
+		return a < b
+	}
+	pq := kpq.NewKeyedPriorityQueue[string](cmp)
+
+	// Insert elements onto the priority queue
+	pq.Push("second", 42)
+	pq.Push("first", 30)
+	pq.Push("last", 50)
+
+	// Updates an element
+	pq.Set("last", 20)
+
+	k, v, ok := pq.Pop()
+	if !ok {
+		log.Fatal("priority queue is empty")
+	}
+
+	fmt.Printf("Key: %q, Value: %d\n", k, v)
+
+	// Inserts a new element
+	pq.Set("new_first", 1)
+	k, v, ok = pq.Pop()
+	if !ok {
+		log.Fatal("priority queue is empty")
+	}
+
+	fmt.Printf("Key: %q, Value: %d\n", k, v)
+	// Output:
+	// Key: "last", Value: 20
+	// Key: "new_first", Value: 1
+}
diff --git a/kpq/keyed_priority_queue.go b/kpq/keyed_priority_queue.go
index 6f99787..1627c45 100644
--- a/kpq/keyed_priority_queue.go
+++ b/kpq/keyed_priority_queue.go
@@ -99,12 +99,16 @@ func (pq *KeyedPriorityQueue[K, V]) Push(k K, v V) error {
 		return newKeyAlreadyExistsError(k)
 	}
 
+	pq.push(k, v)
+	return nil
+}
+
+func (pq *KeyedPriorityQueue[K, V]) push(k K, v V) {
 	n := len(pq.pm)
 	pq.pm = append(pq.pm, k)
 	pq.im[k] = n
 	pq.vals[k] = v
 	pq.swim(n)
-	return nil
 }
 
 // Pop removes and returns the highest priority key and value from the priority queue.
@@ -129,6 +133,20 @@ func (pq *KeyedPriorityQueue[K, V]) Pop() (K, V, bool) {
 	return k, v, true
 }
 
+// Set inserts a new entry in the priority queue with the given key and value,
+// if the key is not present in it; otherwise, it updates the priority value associated with the given key.
+func (pq *KeyedPriorityQueue[K, V]) Set(k K, v V) {
+	pq.mu.Lock()
+	defer pq.mu.Unlock()
+
+	if i, ok := pq.im[k]; ok {
+		pq.update(k, v, i)
+		return
+	}
+
+	pq.push(k, v)
+}
+
 // Update changes the priority value associated with the given key k to the given value v.
 // If there's no key k in the priority queue, it returns a KeyNotFoundError error.
 func (pq *KeyedPriorityQueue[K, V]) Update(k K, v V) error {
@@ -139,10 +157,15 @@ func (pq *KeyedPriorityQueue[K, V]) Update(k K, v V) error {
 	if !ok {
 		return newKeyNotFoundError(k)
 	}
+
+	pq.update(k, v, i)
+	return nil
+}
+
+func (pq *KeyedPriorityQueue[K, V]) update(k K, v V, i int) {
 	pq.vals[k] = v
 	pq.swim(i)
 	pq.sink(i, len(pq.vals))
-	return nil
 }
 
 // Peek returns the highest priority key and value from the priority queue.
diff --git a/kpq/keyed_priority_queue_test.go b/kpq/keyed_priority_queue_test.go
index ba2f48a..9b1b91d 100644
--- a/kpq/keyed_priority_queue_test.go
+++ b/kpq/keyed_priority_queue_test.go
@@ -439,6 +439,49 @@ func TestKeyedPriorityQueue_IsEmpty(t *testing.T) {
 	}
 }
 
+func TestKeyedPriorityQueue_Set(t *testing.T) {
+	pq := NewKeyedPriorityQueue[string](func(x, y int) bool {
+		return x < y
+	})
+
+	items := []struct {
+		key string
+		val int
+	}{
+		{key: "fourth", val: 10},
+		{key: "second", val: 8},
+		{key: "third", val: 9},
+		{key: "first", val: 6},
+		{key: "last", val: 20},
+	}
+
+	for _, item := range items {
+		err := pq.Push(item.key, item.val)
+		if err != nil {
+			t.Fatalf("Push(%v, %v): got unexpected error %v", item.key, item.val, err)
+		}
+	}
+
+	testCases := []struct {
+		k string
+		v int
+	}{
+		{"last", 5},
+		{"new_first", 1},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.k, func(t *testing.T) {
+			pq.Set(tc.k, tc.v)
+
+			if got, _ := pq.PeekValue(); got != tc.v {
+				t.Errorf("pq.PeekValue(): got value %d; want %d", got, tc.v)
+			}
+		})
+	}
+
+}
+
 func benchmarkKeyedPriorityQueue_PushPop(b *testing.B, n int) {
 	pq := NewKeyedPriorityQueue[int](func(a, b int) bool {
 		return a > b