Source file `owl_utils_heap.ml`

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
# 1 "src/base/misc/owl_utils_heap.ml"
(*
 * OWL - OCaml Scientific and Engineering Computing
 * Copyright (c) 2016-2020 Liang Wang <liang.wang@cl.cam.ac.uk>
 *)

(* A naive implementation of a min-heap for Owl's internal use. *)

type 'a t =
  { mutable used : int
  ; mutable data : 'a array
  ; cmp : 'a -> 'a -> int
  }

let left_son pos = (pos lsl 1) + 1

let right_son pos = (pos lsl 1) + 2

let parent pos = (pos - 1) lsr 1

(* can't set an initial size without knowing the type *)
let make cmp = { used = 0; data = [||]; cmp }

let make_int ?(initial_size = 0) cmp = { used = 0; data = Array.make initial_size 0; cmp }

let make_float ?(initial_size = 0) cmp =
  { used = 0; data = Array.make initial_size 0.; cmp }


let size_arr heap = Array.length heap.data

let extend_size heap = heap.data <- Array.(append heap.data (copy heap.data))

let size heap = heap.used

let push ({ used; cmp; _ } as heap) x =
  if size_arr heap = 0
  then (
    heap.data <- [| x |];
    heap.used <- 1)
  else (
    if size_arr heap <= used then extend_size heap;
    (* insert x at the end and swap it with its parent if x is smaller *)
    let pos = ref used in
    let par = ref (parent !pos) in
    while !pos > 0 && cmp x heap.data.(!par) < 0 do
      heap.data.(!pos) <- heap.data.(!par);
      pos := !par;
      par := parent !pos
    done;
    heap.used <- used + 1;
    heap.data.(!pos) <- x)


let pop ({ data; cmp; _ } as heap) =
  match heap.used with
  | 0 -> invalid_arg "owl_utils_heap: can't pop an element from an empty heap"
  | _ ->
    let x = data.(0) in
    (* replace the first element with the last one and swap it with its
     * smallest son *)
    heap.used <- heap.used - 1;
    data.(0) <- data.(heap.used);
    let pos = ref 0 in
    let again = ref true in
    while !again do
      let small = ref !pos in
      let left = left_son !pos
      and right = right_son !pos in
      if left < heap.used && cmp data.(left) data.(!small) < 0 then small := left;
      if right < heap.used && cmp data.(right) data.(!small) < 0 then small := right;
      if !pos = !small
      then again := false
      else (
        let tmp = data.(!pos) in
        data.(!pos) <- data.(!small);
        data.(!small) <- tmp;
        pos := !small)
    done;
    x


let peek heap =
  match heap.used with
  | 0 -> invalid_arg "owl_utils_heap: can't peek at an empty heap"
  | _ -> heap.data.(0)


let is_empty heap = heap.used = 0

let to_array heap = Array.sub heap.data 0 heap.used