basis-library/util/heap.sml

   1 (* Copyright (C) 2007-2007 Wesley W. Terpstra
   2  *
   3  * MLton is released under a BSD-style license.
   4  * See the file MLton-LICENSE for details.
   5  *)
   6
   7 structure Heap:
   8    sig
   9       (* Sorts the provided array relative to the lessthan argument*)
  10       val heapSort: 'a array * ('a * 'a -> bool) -> unit
  11       (* Precondition: array is 0+ true values followed by 0+ false values *)
  12       (* Finds the index of the first array entry where: f x = false *)
  13       val binarySearch: 'a array * ('a -> bool) -> int
  14    end =
  15    struct
  16       fun heapSort (a : 'a array, lessthan : 'a * 'a -> bool) =
  17          let
  18             open Array
  19
  20             (* Push the hole down until value > both children *)
  21             fun pushHoleDown ( hole, end_of_heap, value ) =
  22                let
  23                   val left_child = Int.+ (Int.* (hole, 2), 1)
  24                   val right_child = Int.+ (left_child, 1)
  25                in
  26                   (* Recursion: two children *)
  27                   if Int.< (right_child, end_of_heap)
  28                   then let val left_value = sub (a, left_child)
  29                            val right_value = sub (a, right_child)
  30                            val (bigger_child, bigger_value) =
  31                                if lessthan (left_value, right_value)
  32                                then (right_child, right_value)
  33                                else (left_child, left_value)
  34                        in  if lessthan (bigger_value, value)
  35                            then update (a, hole, value)
  36                            else (update (a, hole, bigger_value);
  37                                  pushHoleDown (bigger_child, end_of_heap, value))
  38                        end
  39                   (* Base case: one child *)
  40                   else if right_child = end_of_heap
  41                   then let val left_value = sub (a, left_child)
  42                        in  if lessthan (left_value, value)
  43                            then update (a, hole, value)
  44                            else (update (a, hole, left_value);
  45                                  update (a, left_child, value))
  46                        end
  47                   (* Base case: no children *)
  48                   else update (a, hole, value)
  49                end
  50
  51             (* Move largest element to end_of_table, then restore invariant *)
  52             fun sortHeap end_of_heap =
  53                let val end_of_heap = Int.- (end_of_heap, 1)
  54                in  if end_of_heap = 0 then () else
  55                    let val value = sub (a, end_of_heap)
  56                    in  update (a, end_of_heap, sub (a, 0));
  57                        pushHoleDown (0, end_of_heap, value);
  58                        sortHeap end_of_heap
  59                end end
  60
  61             (* Start at last node w/ parent, loop till 0: push down *)
  62             val heapSize = Array.length a
  63             fun heapify i =
  64                if i = 0 then () else
  65                let val i = Int.- (i, 1)
  66                in  pushHoleDown (i, heapSize, sub (a, i));
  67                    heapify i
  68                end
  69          in
  70             if Int.<= (heapSize, 1) then () else
  71             (heapify (Int.div (heapSize, 2)); sortHeap heapSize)
  72          end
  73
  74       fun binarySearch (a : 'a array, f : 'a -> bool) =
  75          let
  76             fun loop (lower, upper) =
  77                (* Base case: one element left *)
  78                if Int.- (upper, lower) = 1
  79                then if f (Array.sub (a, lower)) then upper else lower
  80                (* Recursive case: check middle *)
  81                else let val mid = Int.div (Int.+ (lower, upper), 2)
  82                     in  if f (Array.sub (a, mid))
  83                         then loop (mid, upper)
  84                         else loop (lower, mid)
  85                     end
  86             val size = Array.length a
  87          in
  88             if size = 0 then 0 else loop (0, size)
  89          end
  90    end