From 1d1abd5f0f2ba059780ec288b7ded89322cf2f8d Mon Sep 17 00:00:00 2001
From: Rainer Dreyer <rdrey1@uniregistry.com>
Date: Fri, 22 Sep 2017 18:13:21 -0400
Subject: [PATCH] Adds support for custom comparators and `Heap.split/1`

---
 README.md   | 20 ++++++++++++++++++
 lib/heap.ex | 58 ++++++++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 038fd1b..0b89369 100644
--- a/README.md
+++ b/README.md
@@ -56,6 +56,26 @@ Heap.new
 # => [:bread, :milk, :coffee, :eggs, :butter, :jam]
 ```
 
+The heap can also be constructed with a custom comparator:
+
+```elixir
+Heap.new(&(Date.compare(elem(&1, 0), elem(&2, 0)) == :gt))
+|> Heap.push({~D[2017-11-20], :jam})
+|> Heap.push({~D[2017-11-21], :milk})
+|> Heap.push({~D[2017-10-21], :bread})
+|> Heap.push({~D[2017-10-20], :eggs})
+|> Enum.map(fn {_, what} -> what end)
+# => [:milk, :jam, :bread, :eggs]
+```
+
+To access the root and the rest of the heap in one line use `Heap.split/1`:
+
+```elixir
+{root, rest} = Heap.split(heap)
+{root, rest} == {Heap.root(heap), Heap.pop(heap)}
+# => true
+```
+
 ### Documentation
 
 Full API documentation is available on (hexdocs.pm)[https://hexdocs.pm/heap]
diff --git a/lib/heap.ex b/lib/heap.ex
index 5925907..8ccc3fe 100644
--- a/lib/heap.ex
+++ b/lib/heap.ex
@@ -57,19 +57,40 @@ defmodule Heap do
   @doc """
   Create an empty heap with a specific comparator.
 
-  Provide a `comparator` option, which can be either `:<` or `:>` to indicate
-  that the `Heap` should use Elixir's normal `<` or `>` comparison functions.
+  Provide a `comparator` option, which can be `:<`, `:>` to indicate
+  that the `Heap` should use Elixir's normal `<` or `>` comparison functions
+  or a custom comparator function.
 
-  ## Examples
+    ## Examples
 
-      iex> Heap.new(:<)
-      ...>   |> Heap.comparator()
-      :<
+        iex> Heap.new(:<)
+        ...>   |> Heap.comparator()
+        :<
+
+  If given a function it should compare two arguments, and return `true` if
+  the first argument precedes the second one.
+
+    ## Examples
+
+        iex> 1..10
+        ...>   |> Enum.shuffle()
+        ...>   |> Enum.into(Heap.new(&(&1 > &2)))
+        ...>   |> Enum.to_list()
+        [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
+
+        iex> Heap.new(&(Date.compare(elem(&1, 0), elem(&2, 0)) == :gt))
+        ...>   |> Heap.push({~D[2017-11-20], :jam})
+        ...>   |> Heap.push({~D[2017-11-21], :milk})
+        ...>   |> Heap.push({~D[2017-10-21], :bread})
+        ...>   |> Heap.push({~D[2017-10-20], :eggs})
+        ...>   |> Enum.map(fn {_, what} -> what end)
+        [:milk, :jam, :bread, :eggs]
   """
   @spec new(:> | :<) :: t
   def new(:>), do: %Heap{comparator: :>}
   def new(:<), do: %Heap{comparator: :<}
-
+  @spec new(((any, any) -> boolean)) :: t
+  def new(fun), do: %Heap{comparator: fun}
   @doc """
   Test if `heap` is empty.
 
@@ -187,6 +208,19 @@ defmodule Heap do
   @spec comparator(t) :: :< | :>
   def comparator(%Heap{comparator: d}), do: d
 
+  @doc """
+  Return the root element and the rest of the heap in one operation.
+
+  ## Examples
+
+      iex> heap = 1..10 |> Enum.into(Heap.min())
+      ...> rest = Heap.pop(heap)
+      ...> {1, rest} == Heap.split(heap)
+      true
+  """
+  @spec split(t) :: {any, t}
+  def split(%Heap{} = heap), do: {Heap.root(heap), Heap.pop(heap)}
+
   defp meld(nil, queue, _), do: queue
   defp meld(queue, nil, _), do: queue
 
@@ -196,6 +230,13 @@ defmodule Heap do
   defp meld({k0, l0}, {k1, _} = r, :<) when k0 > k1, do: {k0, [r | l0]}
   defp meld({_, _} = l, {k1, r0}, :<), do: {k1, [l | r0]}
 
+  defp meld({k0, l0} = l, {k1, r0} = r, fun) do
+    case fun.(k0, k1) do
+      true -> {k0, [r | l0]}
+      false -> {k1, [l | r0]}
+    end
+  end
+
   defp pair([], _), do: nil
   defp pair([q], _), do: q
   defp pair([q0, q1 | q], d) do
@@ -206,8 +247,7 @@ defmodule Heap do
   defp has_member?(_, previous, compare) when previous == compare, do: true
   defp has_member?(nil, _, _), do: false
   defp has_member?(heap, _, compare) do
-    previous = Heap.root heap
-    heap     = Heap.pop heap
+    {previous, heap} = Heap.split heap
     has_member? heap, previous, compare
   end
 end