Tuesday, February 3, 2015

Implement max heap in Ocaml


(* In this script, we will implement the max heap.
 * In this implementation, we use the first element in
 * the array to store the number of elements in the heap.
 * This number has type of int, it follows that the heap
 * is of type int array.
 *)


let parent i = i / 2 ;;
let left i = 2 * i ;;
let right i = 2 * i + 1 ;;


(* the function heapify maintains the heap property of the sub-heap
 * rooted at k
 *)

let rec heapify heap k = 
  let count = heap.(0) in 
  if k < 1 then
    raise (Invalid_argument "heap")
  else if k < count then
    match left k <= count, right k <= count with
    |(true, true) ->
        let elem = heap.(k) 
        and elem_left = heap.(left k) 
        and elem_right = heap.(right k) in
        (match elem < elem_left, elem < elem_right with
        | (false, false ) -> ()
        | (_,_) ->
            if elem_left < elem_right then
                (heap.(k) <- elem_right; heap.(right k) <- elem; heapify heap (right k))
            else
                (heap.(k) <- elem_left; heap.(left k) <- elem; heapify heap (left k)))
    |(true, false) -> 
        let elem_left = heap.(left k) and elem = heap.(k) in
        if elem <elem_left then
          (heap.(left k) <- elem ; heap.(k) <- elem_left)
    |(false, _)  -> () ;;



(* the function heapify_all will call function heapify on every node in the heap *)

let heapify_all  heap = 
  let count = heap.(0) in
  for i = count downto 1 do
    heapify heap i
  done;;

(* add function will add a new element to the heap *)
let add heap elem = 
  let count = heap.(0) in
  if count = Array.length heap then
    raise (Failure "The heap is full. Cannot add new element.")
  else
    (heap.(0) <- count + 1;
     let p = ref heap.(0) in
     while !p / 2  > 0 && heap.(!p) > heap.(!p / 2)  do
       p := !p / 2;
       heapify heap !p
     done);;


(* max_heap return the maximum value of the heap but not remove it from the heap
 * *)
let max_heap heap =
  if heap.(0) = 0 then
    raise (Failure "The heap is empty.")
  else
    heap.(1) ;;

(* pop function returns the maximum value of the heap and remove it from the
 * heap *)

let pop heap = 
  let count = heap.(0) in
  if count  = 0 then
    raise (Failure "The heap is empty.")
  else
    let elem = heap.(1) in    
    heap.(1) <- heap.(count);
    heap.(0) <- count - 1;
    heapify heap 1;
    elem;;
   

No comments:

Post a Comment