github.com/jxskiss/gopkg/v2@v2.14.9-0.20240514120614-899f3e7952b4/collection/heapx/priority_queue_test.go (about)

     1  package heapx
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/assert"
     7  )
     8  
     9  func TestPriorityQueue(t *testing.T) {
    10  	nums := []int64{2, 0, 1, 5, 9, 6, 4, 7, 8, 3}
    11  
    12  	t.Run("min queue", func(t *testing.T) {
    13  		pq := NewMinPriorityQueue[int64, *int64]()
    14  		for i := range nums {
    15  			num := nums[i]
    16  			pq.Push(num, &num)
    17  		}
    18  		assert.True(t, pq.Len() == 10)
    19  
    20  		for i := range nums {
    21  			priority, value, ok := pq.Peek()
    22  			assert.True(t, ok)
    23  			assert.Equal(t, int64(i), *value)
    24  			assert.Equal(t, int64(i), priority)
    25  
    26  			priority, value, ok = pq.Pop()
    27  			assert.True(t, ok)
    28  			assert.Equal(t, int64(i), *value)
    29  			assert.Equal(t, int64(i), priority)
    30  		}
    31  
    32  		assert.Equal(t, 1, len(pq.heap.items.ss))
    33  		assert.Equal(t, initSize, pq.heap.items.cap)
    34  		assert.Equal(t, 0, pq.heap.items.len)
    35  	})
    36  
    37  	t.Run("max queue", func(t *testing.T) {
    38  		pq := NewMaxPriorityQueue[int64, *int64]()
    39  		for i := range nums {
    40  			num := nums[i]
    41  			pq.Push(num, &num)
    42  		}
    43  		assert.True(t, pq.Len() == 10)
    44  
    45  		for i := range nums {
    46  			priority, value, ok := pq.Peek()
    47  			assert.True(t, ok)
    48  			assert.Equal(t, int64(9-i), *value)
    49  			assert.Equal(t, int64(9-i), priority)
    50  
    51  			priority, value, ok = pq.Pop()
    52  			assert.True(t, ok)
    53  			assert.Equal(t, int64(9-i), *value)
    54  			assert.Equal(t, int64(9-i), priority)
    55  		}
    56  
    57  		assert.Equal(t, 1, len(pq.heap.items.ss))
    58  		assert.Equal(t, initSize, pq.heap.items.cap)
    59  		assert.Equal(t, 0, pq.heap.items.len)
    60  	})
    61  }