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

     1  package heapx
     2  
     3  import (
     4  	"math/rand"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  )
     9  
    10  func TestHeap(t *testing.T) {
    11  	nums := make([]int, 3333)
    12  	for i := range nums {
    13  		nums[i] = i
    14  	}
    15  	rand.Shuffle(len(nums), func(i, j int) {
    16  		nums[i], nums[j] = nums[j], nums[i]
    17  	})
    18  
    19  	t.Run("min heap", func(t *testing.T) {
    20  		h1 := NewHeap[int](func(lhs, rhs int) bool {
    21  			return lhs < rhs
    22  		})
    23  		assert.True(t, h1.Len() == 0)
    24  		for i := range nums {
    25  			h1.Push(nums[i])
    26  		}
    27  		for i := range nums {
    28  			x, ok := h1.Pop()
    29  			assert.True(t, ok)
    30  			assert.Equal(t, i, x)
    31  		}
    32  		x, ok := h1.Peek()
    33  		assert.False(t, ok)
    34  		assert.Equal(t, 0, x)
    35  		assert.Equal(t, 1, len(h1.items.ss))
    36  		assert.Equal(t, bktSize, h1.items.cap)
    37  		assert.Equal(t, 0, h1.items.len)
    38  	})
    39  
    40  	t.Run("max heap", func(t *testing.T) {
    41  		h2 := NewHeap[int](func(lhs, rhs int) bool {
    42  			return rhs < lhs
    43  		})
    44  		assert.True(t, h2.Len() == 0)
    45  		for i := range nums {
    46  			h2.Push(nums[i])
    47  		}
    48  		for i := len(nums) - 1; i >= 0; i-- {
    49  			x, ok := h2.Pop()
    50  			assert.True(t, ok)
    51  			assert.Equal(t, i, x)
    52  		}
    53  		x, ok := h2.Peek()
    54  		assert.False(t, ok)
    55  		assert.Equal(t, 0, x)
    56  		assert.Equal(t, 1, len(h2.items.ss))
    57  		assert.Equal(t, bktSize, h2.items.cap)
    58  		assert.Equal(t, 0, h2.items.len)
    59  	})
    60  
    61  	t.Run("shrink", func(t *testing.T) {
    62  		h3 := NewHeap[int](func(lhs, rhs int) bool {
    63  			return lhs < rhs
    64  		})
    65  		assert.Equal(t, 0, h3.Len())
    66  		assert.Equal(t, 0, h3.items.cap)
    67  		assert.Equal(t, 0, len(h3.items.ss))
    68  
    69  		for i := 0; i < bktSize; i++ {
    70  			h3.Push(i)
    71  		}
    72  		assert.Equal(t, bktSize, h3.Len())
    73  		assert.Equal(t, bktSize, h3.items.cap)
    74  		assert.Equal(t, 1, len(h3.items.ss))
    75  
    76  		for i := 0; i < bktSize; i++ {
    77  			h3.Push(i)
    78  		}
    79  		assert.Equal(t, 2*bktSize, h3.Len())
    80  		assert.Equal(t, 2*bktSize, h3.items.cap)
    81  		assert.Equal(t, 2, len(h3.items.ss))
    82  
    83  		for i := 0; i < 10; i++ {
    84  			h3.Push(i)
    85  		}
    86  		assert.Equal(t, 2*bktSize+10, h3.Len())
    87  		assert.Equal(t, 3*bktSize, h3.items.cap)
    88  		assert.Equal(t, 3, len(h3.items.ss))
    89  
    90  		for i := 0; i < 10; i++ {
    91  			h3.Pop()
    92  		}
    93  		assert.Equal(t, 2*bktSize, h3.Len())
    94  		assert.Equal(t, 3*bktSize, h3.items.cap)
    95  		assert.Equal(t, 3, len(h3.items.ss))
    96  
    97  		for i := 0; i < shrinkThreshold-1; i++ {
    98  			h3.Pop()
    99  		}
   100  		assert.Equal(t, 2*bktSize-(shrinkThreshold-1), h3.Len())
   101  		assert.Equal(t, 3*bktSize, h3.items.cap)
   102  		assert.Equal(t, 3, len(h3.items.ss))
   103  
   104  		h3.Pop()
   105  		assert.Equal(t, 2*bktSize-shrinkThreshold, h3.Len())
   106  		assert.Equal(t, 2*bktSize, h3.items.cap)
   107  		assert.Equal(t, 2, len(h3.items.ss))
   108  
   109  		for i := 0; i < (bktSize - shrinkThreshold); i++ {
   110  			h3.Pop()
   111  		}
   112  		assert.Equal(t, bktSize, h3.Len())
   113  		assert.Equal(t, 2*bktSize, h3.items.cap)
   114  		assert.Equal(t, 2, len(h3.items.ss))
   115  
   116  		for i := 0; i < shrinkThreshold-1; i++ {
   117  			h3.Pop()
   118  		}
   119  		assert.Equal(t, bktSize-(shrinkThreshold-1), h3.Len())
   120  		assert.Equal(t, 2*bktSize, h3.items.cap)
   121  		assert.Equal(t, 2, len(h3.items.ss))
   122  
   123  		h3.Pop()
   124  		assert.Equal(t, bktSize-shrinkThreshold, h3.Len())
   125  		assert.Equal(t, bktSize, h3.items.cap)
   126  		assert.Equal(t, 1, len(h3.items.ss))
   127  
   128  		for i := 0; i < bktSize-shrinkThreshold; i++ {
   129  			h3.Pop()
   130  		}
   131  		assert.Equal(t, 0, h3.Len())
   132  		assert.Equal(t, bktSize, h3.items.cap)
   133  		assert.Equal(t, 1, len(h3.items.ss))
   134  	})
   135  }