Golang heap.Interface类代码示例

func siftUpMin(h heap.Interface, i int) {
	for gp := parent(parent(i)); gp >= 0; i, gp = gp, parent(parent(gp)) {
		if h.Less(gp, i) {
			break
		}
		h.Swap(i, gp)
	}
}

func minGrandchild(h heap.Interface, i, n int) int {
	cg := (i*2+1)*2 + 1
	m := cg
	for i := 1; i <= 3; i++ {
		if h.Less(cg+i, m) {
			m = cg + i
		}
	}
	return m
}

// Index of minimum of children and grandchildren (if any) of i.
// Precondition: i has at least one child.
func minTwoGen(h heap.Interface, i, n int) int {
	c1, c2 := i*2+1, i*2+2
	m := c1
	for _, k := range []int{c2, c1*2 + 1, c1*2 + 2, c2*2 + 1, c2*2 + 2} {
		if k >= n {
			break
		}
		if h.Less(k, m) {
			m = k
		}
	}
	return m
}

func TestBasicDBHeapBehavior(t *testing.T) {
	var dbheap heap.Interface

	testutil.VerifyTestType(t, testutil.UnitTestType)

	Convey("With an empty dbHeap", t, func() {
		dbheap = &DBHeap{}
		heap.Init(dbheap)

		Convey("when inserting unordered dbCounters with different active counts", func() {
			heap.Push(dbheap, &dbCounter{75, nil})
			heap.Push(dbheap, &dbCounter{121, nil})
			heap.Push(dbheap, &dbCounter{76, nil})
			heap.Push(dbheap, &dbCounter{51, nil})
			heap.Push(dbheap, &dbCounter{82, nil})
			heap.Push(dbheap, &dbCounter{117, nil})
			heap.Push(dbheap, &dbCounter{49, nil})
			heap.Push(dbheap, &dbCounter{101, nil})
			heap.Push(dbheap, &dbCounter{122, nil})
			heap.Push(dbheap, &dbCounter{33, nil})
			heap.Push(dbheap, &dbCounter{0, nil})

			Convey("they should pop in active order, least to greatest", func() {
				prev := -1
				for dbheap.Len() > 0 {
					popped := heap.Pop(dbheap).(*dbCounter)
					So(popped.active, ShouldBeGreaterThan, prev)
					prev = popped.active
				}
			})
		})

		Convey("when inserting unordered dbCounters with different bson sizes", func() {
			heap.Push(dbheap, &dbCounter{0, []*Intent{&Intent{Size: 70}}})
			heap.Push(dbheap, &dbCounter{0, []*Intent{&Intent{Size: 1024}}})
			heap.Push(dbheap, &dbCounter{0, []*Intent{&Intent{Size: 97}}})
			heap.Push(dbheap, &dbCounter{0, []*Intent{&Intent{Size: 3}}})
			heap.Push(dbheap, &dbCounter{0, []*Intent{&Intent{Size: 1024 * 1024}}})

			Convey("they should pop in bson size order, greatest to least", func() {
				prev := int64(1024*1024 + 1) // Maximum
				for dbheap.Len() > 0 {
					popped := heap.Pop(dbheap).(*dbCounter)
					So(popped.collections[0].Size, ShouldBeLessThan, prev)
					prev = popped.collections[0].Size
				}
			})
		})
	})
}

// Index of maximum element in h.
//
// There's no MinInd because the minimum of a heap is always at index zero.
func MaxInd(h heap.Interface) (maxind int) {
	switch h.Len() {
	case 1:
		maxind = 0
	case 2:
		maxind = 1
	default:
		maxind = 1
		if h.Less(1, 2) {
			maxind = 2
		}
	}
	return
}

// Removes and returns the maximum element of h.
func PopMax(h heap.Interface) interface{} {
	n := h.Len()
	if n <= 2 {
		return h.Pop()
	}

	i := 1
	if h.Less(1, 2) {
		i = 2
	}
	h.Swap(i, n-1)
	x := h.Pop()
	siftDownMax(h, i, n-1)
	return x
}

// Removes and returns the minimum element of h.
func PopMin(h heap.Interface) interface{} {
	n := h.Len() - 1
	h.Swap(0, n)
	x := h.Pop()
	siftDownMin(h, 0, n)
	return x
}

func siftUp(h heap.Interface, i int) {
	if minLevel(i) {
		if i > 0 && h.Less(parent(i), i) {
			h.Swap(i, parent(i))
			siftUpMax(h, parent(i))
		} else {
			siftUpMin(h, i)
		}
	} else {
		if i > 0 && h.Less(i, parent(i)) {
			h.Swap(i, parent(i))
			siftUpMin(h, parent(i))
		} else {
			siftUpMax(h, i)
		}
	}
}

// Adds x to the heap x.
func Push(h heap.Interface, x interface{}) {
	h.Push(x)
	siftUp(h, h.Len()-1)
}

// Establishes heap order in h.
//
// Every non-empty heap must be initialized using this function before any of
// the other functions of this package is used on it.
//
// Note that, although this package uses the interface from container/heap,
// you must not mix operations from both packages.
func Init(h heap.Interface) {
	n := h.Len()
	for i := parent(n); i >= 0; i-- {
		siftDown(h, i, n)
	}
}

func siftDownMin(h heap.Interface, i, n int) {
	for (i*2+2)*2+2 < n { // has four grandchildren
		m := minGrandchild(h, i, n)
		if h.Less(i, m) {
			return
		}
		h.Swap(i, m)
		if h.Less(parent(m), m) {
			h.Swap(m, parent(m))
		}
		i = m
	}

	if hasChild(i, n) {
		m := minTwoGen(h, i, n)
		if m > i*2+2 { // must be a grandchild
			if h.Less(m, i) {
				h.Swap(i, m)
				if h.Less(parent(m), m) {
					h.Swap(m, parent(m))
				}
				siftDownMin(h, m, n)
			}
		} else if h.Less(m, i) {
			h.Swap(i, m)
		}
	}
}

// Iterative version after Bojesen, "Heap implementations and variations",
// http://www.diku.dk/forskning/performance-engineering/Jesper/heaplab/heapsurvey_html/node11.html
func siftDownMax(h heap.Interface, i, n int) {
	for (i*2+2)*2+2 < n { // has four grandchildren
		m := maxGrandchild(h, i, n)
		if h.Less(m, i) {
			return
		}
		h.Swap(i, m)
		if h.Less(m, parent(m)) {
			h.Swap(m, parent(m))
		}
		i = m
	}

	// XXX Following is from the original paper; couldn't get Bojesen's version
	// to work.
	if hasChild(i, n) {
		m := maxTwoGen(h, i, n)
		if m > i*2+2 { // must be a grandchild
			if h.Less(i, m) {
				h.Swap(i, m)
				if h.Less(m, parent(m)) {
					h.Swap(m, parent(m))
				}
				siftDownMax(h, m, n)
			}
		} else if h.Less(i, m) {
			h.Swap(i, m)
		}
	}
}