heap.pony

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
type MinHeap[A: Comparable[A] #read] is BinaryHeap[A, MinHeapPriority[A]]
type MaxHeap[A: Comparable[A] #read] is BinaryHeap[A, MaxHeapPriority[A]]

class BinaryHeap[A: Comparable[A] #read, P: BinaryHeapPriority[A]]
  """
  A priority queue implemented as a binary heap. The `BinaryHeapPriority` type
  parameter determines whether this is max-heap or a min-heap.
  """
  embed _data: Array[A]

  new create(len: USize) =>
    """
    Create an empty heap with space for `len` elements.
    """
    _data = Array[A](len)

  fun ref clear() =>
    """
    Remove all elements from the heap.
    """
    _data.clear()

  fun size(): USize =>
    """
    Return the number of elements in the heap.
    """
    _data.size()

  fun peek(): this->A ? =>
    """
    Return the highest priority item in the heap. For max-heaps, the greatest
    item will be returned. For min-heaps, the smallest item will be returned.
    """
    _data(0)?

  fun ref push(value: A) =>
    """
    Push an item into the heap.

    The time complexity of this operation is O(log(n)) with respect to the size
    of the heap.
    """
    _data.push(value)
    _sift_up(size() - 1)

  fun ref pop(): A^ ? =>
    """
    Remove the highest priority value from the heap and return it. For
    max-heaps, the greatest item will be returned. For min-heaps, the smallest
    item will be returned.

    The time complexity of this operation is O(log(n)) with respect to the size
    of the heap.
    """
    let n = size() - 1
    _data.swap_elements(0, n)?
    _sift_down(0, n)
    _data.pop()?

  fun ref append(
    seq: (ReadSeq[A] & ReadElement[A^]),
    offset: USize = 0,
    len: USize = -1)
  =>
    """
    Append len elements from a sequence, starting from the given offset.
    """
    _data.append(seq, offset, len)
    _make_heap()

  fun ref concat(iter: Iterator[A^], offset: USize = 0, len: USize = -1) =>
    """
    Add len iterated elements, starting from the given offset.
    """
    _data.concat(iter, offset, len)
    _make_heap()

  fun values(): ArrayValues[A, this->Array[A]]^ =>
    """
    Return an iterator for the elements in the heap. The order of elements is
    arbitrary.
    """
    _data.values()

  fun ref _make_heap() =>
    let n = size()
    if n < 2 then return end
    var i = (n / 2)
    while (i = i - 1) > 0 do
      _sift_down(i, n)
    end

  fun ref _sift_up(n: USize) =>
    var idx = n
    try
      while true do
        let parent_idx = (idx - 1) / 2
        if (parent_idx == idx) or not P(_data(idx)?, _data(parent_idx)?) then
          break
        end
        _data.swap_elements(parent_idx, idx)?
        idx = parent_idx
      end
    end

  fun ref _sift_down(start: USize, n: USize): Bool =>
    var idx = start
    try
      while true do
        var left = (2 * idx) + 1
        if (left >= n) or (left < 0) then
          break
        end
        let right = left + 1
        if (right < n) and P(_data(right)?, _data(left)?) then
          left = right
        end
        if not P(_data(left)?, _data(idx)?) then
          break
        end
        _data.swap_elements(idx, left)?
        idx = left
      end
    end
    idx > start

  fun _apply(i: USize): this->A ? =>
    _data(i)?

type BinaryHeapPriority[A: Comparable[A] #read] is
  ( _BinaryHeapPriority[A]
  & (MinHeapPriority[A] | MaxHeapPriority[A]))

interface val _BinaryHeapPriority[A: Comparable[A] #read]
  new val create()
  fun apply(x: A, y: A): Bool

primitive MinHeapPriority[A: Comparable[A] #read] is _BinaryHeapPriority[A]
  fun apply(x: A, y: A): Bool =>
    x < y

primitive MaxHeapPriority [A: Comparable[A] #read] is _BinaryHeapPriority[A]
  fun apply(x: A, y: A): Bool =>
    x > y