1// Copyright (c) 2015, Emir Pasic. 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 binaryheap implements a binary heap backed by array list.
  6//
  7// Comparator defines this heap as either min or max heap.
  8//
  9// Structure is not thread safe.
 10//
 11// References: http://en.wikipedia.org/wiki/Binary_heap
 12package binaryheap
 13
 14import (
 15	"fmt"
 16	"github.com/emirpasic/gods/lists/arraylist"
 17	"github.com/emirpasic/gods/trees"
 18	"github.com/emirpasic/gods/utils"
 19	"strings"
 20)
 21
 22// Assert Tree implementation
 23var _ trees.Tree = (*Heap)(nil)
 24
 25// Heap holds elements in an array-list
 26type Heap struct {
 27	list       *arraylist.List
 28	Comparator utils.Comparator
 29}
 30
 31// NewWith instantiates a new empty heap tree with the custom comparator.
 32func NewWith(comparator utils.Comparator) *Heap {
 33	return &Heap{list: arraylist.New(), Comparator: comparator}
 34}
 35
 36// NewWithIntComparator instantiates a new empty heap with the IntComparator, i.e. elements are of type int.
 37func NewWithIntComparator() *Heap {
 38	return &Heap{list: arraylist.New(), Comparator: utils.IntComparator}
 39}
 40
 41// NewWithStringComparator instantiates a new empty heap with the StringComparator, i.e. elements are of type string.
 42func NewWithStringComparator() *Heap {
 43	return &Heap{list: arraylist.New(), Comparator: utils.StringComparator}
 44}
 45
 46// Push adds a value onto the heap and bubbles it up accordingly.
 47func (heap *Heap) Push(values ...interface{}) {
 48	if len(values) == 1 {
 49		heap.list.Add(values[0])
 50		heap.bubbleUp()
 51	} else {
 52		// Reference: https://en.wikipedia.org/wiki/Binary_heap#Building_a_heap
 53		for _, value := range values {
 54			heap.list.Add(value)
 55		}
 56		size := heap.list.Size()/2 + 1
 57		for i := size; i >= 0; i-- {
 58			heap.bubbleDownIndex(i)
 59		}
 60	}
 61}
 62
 63// Pop removes top element on heap and returns it, or nil if heap is empty.
 64// Second return parameter is true, unless the heap was empty and there was nothing to pop.
 65func (heap *Heap) Pop() (value interface{}, ok bool) {
 66	value, ok = heap.list.Get(0)
 67	if !ok {
 68		return
 69	}
 70	lastIndex := heap.list.Size() - 1
 71	heap.list.Swap(0, lastIndex)
 72	heap.list.Remove(lastIndex)
 73	heap.bubbleDown()
 74	return
 75}
 76
 77// Peek returns top element on the heap without removing it, or nil if heap is empty.
 78// Second return parameter is true, unless the heap was empty and there was nothing to peek.
 79func (heap *Heap) Peek() (value interface{}, ok bool) {
 80	return heap.list.Get(0)
 81}
 82
 83// Empty returns true if heap does not contain any elements.
 84func (heap *Heap) Empty() bool {
 85	return heap.list.Empty()
 86}
 87
 88// Size returns number of elements within the heap.
 89func (heap *Heap) Size() int {
 90	return heap.list.Size()
 91}
 92
 93// Clear removes all elements from the heap.
 94func (heap *Heap) Clear() {
 95	heap.list.Clear()
 96}
 97
 98// Values returns all elements in the heap.
 99func (heap *Heap) Values() []interface{} {
100	values := make([]interface{}, heap.list.Size(), heap.list.Size())
101	for it := heap.Iterator(); it.Next(); {
102		values[it.Index()] = it.Value()
103	}
104	return values
105}
106
107// String returns a string representation of container
108func (heap *Heap) String() string {
109	str := "BinaryHeap\n"
110	values := []string{}
111	for it := heap.Iterator(); it.Next(); {
112		values = append(values, fmt.Sprintf("%v", it.Value()))
113	}
114	str += strings.Join(values, ", ")
115	return str
116}
117
118// Performs the "bubble down" operation. This is to place the element that is at the root
119// of the heap in its correct place so that the heap maintains the min/max-heap order property.
120func (heap *Heap) bubbleDown() {
121	heap.bubbleDownIndex(0)
122}
123
124// Performs the "bubble down" operation. This is to place the element that is at the index
125// of the heap in its correct place so that the heap maintains the min/max-heap order property.
126func (heap *Heap) bubbleDownIndex(index int) {
127	size := heap.list.Size()
128	for leftIndex := index<<1 + 1; leftIndex < size; leftIndex = index<<1 + 1 {
129		rightIndex := index<<1 + 2
130		smallerIndex := leftIndex
131		leftValue, _ := heap.list.Get(leftIndex)
132		rightValue, _ := heap.list.Get(rightIndex)
133		if rightIndex < size && heap.Comparator(leftValue, rightValue) > 0 {
134			smallerIndex = rightIndex
135		}
136		indexValue, _ := heap.list.Get(index)
137		smallerValue, _ := heap.list.Get(smallerIndex)
138		if heap.Comparator(indexValue, smallerValue) > 0 {
139			heap.list.Swap(index, smallerIndex)
140		} else {
141			break
142		}
143		index = smallerIndex
144	}
145}
146
147// Performs the "bubble up" operation. This is to place a newly inserted
148// element (i.e. last element in the list) in its correct place so that
149// the heap maintains the min/max-heap order property.
150func (heap *Heap) bubbleUp() {
151	index := heap.list.Size() - 1
152	for parentIndex := (index - 1) >> 1; index > 0; parentIndex = (index - 1) >> 1 {
153		indexValue, _ := heap.list.Get(index)
154		parentValue, _ := heap.list.Get(parentIndex)
155		if heap.Comparator(parentValue, indexValue) <= 0 {
156			break
157		}
158		heap.list.Swap(index, parentIndex)
159		index = parentIndex
160	}
161}
162
163// Check that the index is within bounds of the list
164func (heap *Heap) withinRange(index int) bool {
165	return index >= 0 && index < heap.list.Size()
166}