diff --git a/core/src/main/scala/cats/collections/Heap.scala b/core/src/main/scala/cats/collections/Heap.scala index 98a71dc3..f6078442 100644 --- a/core/src/main/scala/cats/collections/Heap.scala +++ b/core/src/main/scala/cats/collections/Heap.scala @@ -129,6 +129,19 @@ sealed abstract class Heap[A] { def heapify(a: List[A])(implicit order: Order[A]): Heap[A] = Heap.heapify(a) + def toIterator(implicit order: Order[A]): Iterator[A] = { + @scala.annotation.tailrec + def build(heap: Heap[A], acc: Stream[A]): Stream[A] = heap match { + case Leaf() => acc + case b@Branch(_, _, _) => b.pop match { + case Some((m, h)) => build(h, acc.append(Seq(m))) + case None => acc + } + } + + build(this, Stream.empty).toIterator + } + /** * Remove the min element from the heap (the root) and return it along with the updated heap. * Order O(log n) @@ -171,16 +184,7 @@ sealed abstract class Heap[A] { /** * Returns a sorted list of the elements within the heap. */ - def toList(implicit order: Order[A]): List[A] = { - @tailrec - def loop(h: Heap[A], acc: List[A]): List[A] = - h match { - case Branch(m, _, _) => loop(h.remove, m :: acc) - case Leaf() => acc.reverse - } - - loop(this, Nil) - } + def toList(implicit order: Order[A]): List[A] = toIterator.toList /** * do a foldLeft in the same order as toList. diff --git a/tests/src/test/scala/cats/collections/HeapSpec.scala b/tests/src/test/scala/cats/collections/HeapSpec.scala index c3bee565..a40142a4 100644 --- a/tests/src/test/scala/cats/collections/HeapSpec.scala +++ b/tests/src/test/scala/cats/collections/HeapSpec.scala @@ -271,4 +271,35 @@ class HeapSpec extends CatsSuite { law(h, h) } } + + test("toIterator maintains order")(forAll { heap: Heap[Int] => + val numbers = heap.toIterator.toList + + numbers should contain inOrderElementsOf (numbers.sorted) + }) + + test("toIterator maintains the inverse order")(forAll { xs: List[Int] => + whenever(xs.size > 1) { + val order = Order.from[Int]((a, b) => a.compareTo(b) * -1) + val inverseOrder = Order.from[Int]((a, b) => a.compareTo(b)) + + val heap = Heap.fromIterable(xs)(order) + + val it = heap.toIterator(order) + var min = it.next() + + while (it.hasNext) { + val current = it.next() + + order.gteqv(current, min) should be (true) + inverseOrder.lteqv(current, min) should be (true) + + min = current + } + } + }) + + test("toIterator should match with toList")(forAll {heap: Heap[Int] => + heap.toIterator.toList === heap.toList + }) }