...

Source file src/container/heap/heap_test.go

Documentation: container/heap

     1  // Copyright 2009 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package heap
     6  
     7  import (
     8  	"math/rand"
     9  	"testing"
    10  )
    11  
    12  type myHeap []int
    13  
    14  func (h *myHeap) Less(i, j int) bool {
    15  	return (*h)[i] < (*h)[j]
    16  }
    17  
    18  func (h *myHeap) Swap(i, j int) {
    19  	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
    20  }
    21  
    22  func (h *myHeap) Len() int {
    23  	return len(*h)
    24  }
    25  
    26  func (h *myHeap) Pop() (v any) {
    27  	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
    28  	return
    29  }
    30  
    31  func (h *myHeap) Push(v any) {
    32  	*h = append(*h, v.(int))
    33  }
    34  
    35  func (h myHeap) verify(t *testing.T, i int) {
    36  	t.Helper()
    37  	n := h.Len()
    38  	j1 := 2*i + 1
    39  	j2 := 2*i + 2
    40  	if j1 < n {
    41  		if h.Less(j1, i) {
    42  			t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h[i], j1, h[j1])
    43  			return
    44  		}
    45  		h.verify(t, j1)
    46  	}
    47  	if j2 < n {
    48  		if h.Less(j2, i) {
    49  			t.Errorf("heap invariant invalidated [%d] = %d > [%d] = %d", i, h[i], j1, h[j2])
    50  			return
    51  		}
    52  		h.verify(t, j2)
    53  	}
    54  }
    55  
    56  func TestInit0(t *testing.T) {
    57  	h := new(myHeap)
    58  	for i := 20; i > 0; i-- {
    59  		h.Push(0) // all elements are the same
    60  	}
    61  	Init(h)
    62  	h.verify(t, 0)
    63  
    64  	for i := 1; h.Len() > 0; i++ {
    65  		x := Pop(h).(int)
    66  		h.verify(t, 0)
    67  		if x != 0 {
    68  			t.Errorf("%d.th pop got %d; want %d", i, x, 0)
    69  		}
    70  	}
    71  }
    72  
    73  func TestInit1(t *testing.T) {
    74  	h := new(myHeap)
    75  	for i := 20; i > 0; i-- {
    76  		h.Push(i) // all elements are different
    77  	}
    78  	Init(h)
    79  	h.verify(t, 0)
    80  
    81  	for i := 1; h.Len() > 0; i++ {
    82  		x := Pop(h).(int)
    83  		h.verify(t, 0)
    84  		if x != i {
    85  			t.Errorf("%d.th pop got %d; want %d", i, x, i)
    86  		}
    87  	}
    88  }
    89  
    90  func Test(t *testing.T) {
    91  	h := new(myHeap)
    92  	h.verify(t, 0)
    93  
    94  	for i := 20; i > 10; i-- {
    95  		h.Push(i)
    96  	}
    97  	Init(h)
    98  	h.verify(t, 0)
    99  
   100  	for i := 10; i > 0; i-- {
   101  		Push(h, i)
   102  		h.verify(t, 0)
   103  	}
   104  
   105  	for i := 1; h.Len() > 0; i++ {
   106  		x := Pop(h).(int)
   107  		if i < 20 {
   108  			Push(h, 20+i)
   109  		}
   110  		h.verify(t, 0)
   111  		if x != i {
   112  			t.Errorf("%d.th pop got %d; want %d", i, x, i)
   113  		}
   114  	}
   115  }
   116  
   117  func TestRemove0(t *testing.T) {
   118  	h := new(myHeap)
   119  	for i := 0; i < 10; i++ {
   120  		h.Push(i)
   121  	}
   122  	h.verify(t, 0)
   123  
   124  	for h.Len() > 0 {
   125  		i := h.Len() - 1
   126  		x := Remove(h, i).(int)
   127  		if x != i {
   128  			t.Errorf("Remove(%d) got %d; want %d", i, x, i)
   129  		}
   130  		h.verify(t, 0)
   131  	}
   132  }
   133  
   134  func TestRemove1(t *testing.T) {
   135  	h := new(myHeap)
   136  	for i := 0; i < 10; i++ {
   137  		h.Push(i)
   138  	}
   139  	h.verify(t, 0)
   140  
   141  	for i := 0; h.Len() > 0; i++ {
   142  		x := Remove(h, 0).(int)
   143  		if x != i {
   144  			t.Errorf("Remove(0) got %d; want %d", x, i)
   145  		}
   146  		h.verify(t, 0)
   147  	}
   148  }
   149  
   150  func TestRemove2(t *testing.T) {
   151  	N := 10
   152  
   153  	h := new(myHeap)
   154  	for i := 0; i < N; i++ {
   155  		h.Push(i)
   156  	}
   157  	h.verify(t, 0)
   158  
   159  	m := make(map[int]bool)
   160  	for h.Len() > 0 {
   161  		m[Remove(h, (h.Len()-1)/2).(int)] = true
   162  		h.verify(t, 0)
   163  	}
   164  
   165  	if len(m) != N {
   166  		t.Errorf("len(m) = %d; want %d", len(m), N)
   167  	}
   168  	for i := 0; i < len(m); i++ {
   169  		if !m[i] {
   170  			t.Errorf("m[%d] doesn't exist", i)
   171  		}
   172  	}
   173  }
   174  
   175  func BenchmarkDup(b *testing.B) {
   176  	const n = 10000
   177  	h := make(myHeap, 0, n)
   178  	for i := 0; i < b.N; i++ {
   179  		for j := 0; j < n; j++ {
   180  			Push(&h, 0) // all elements are the same
   181  		}
   182  		for h.Len() > 0 {
   183  			Pop(&h)
   184  		}
   185  	}
   186  }
   187  
   188  func TestFix(t *testing.T) {
   189  	h := new(myHeap)
   190  	h.verify(t, 0)
   191  
   192  	for i := 200; i > 0; i -= 10 {
   193  		Push(h, i)
   194  	}
   195  	h.verify(t, 0)
   196  
   197  	if (*h)[0] != 10 {
   198  		t.Fatalf("Expected head to be 10, was %d", (*h)[0])
   199  	}
   200  	(*h)[0] = 210
   201  	Fix(h, 0)
   202  	h.verify(t, 0)
   203  
   204  	for i := 100; i > 0; i-- {
   205  		elem := rand.Intn(h.Len())
   206  		if i&1 == 0 {
   207  			(*h)[elem] *= 2
   208  		} else {
   209  			(*h)[elem] /= 2
   210  		}
   211  		Fix(h, elem)
   212  		h.verify(t, 0)
   213  	}
   214  }
   215  

View as plain text