exercises/ch01/tree.ml

   1 module Array = ArrayLabels
   2 module List = ListLabels
   3
   4 module type TREE = sig
   5   type ('k, 'v) t
   6
   7   val empty : ('k, 'v) t
   8
   9   val set : ('k, 'v) t -> k:'k -> v:'v -> ('k, 'v) t
  10
  11   val get : ('k, 'v) t -> k:'k -> 'v option
  12
  13   val member : ('k, 'v) t -> k:'k -> bool
  14
  15   val print
  16     : ('k, 'v) t
  17     -> k_to_string:('k -> string)
  18     -> indent_level:string
  19     -> unit
  20 end
  21
  22 module BinaryTree : TREE = struct
  23   type ('k, 'v) t =
  24     | Node of 'k * 'v * ('k, 'v) t * ('k, 'v) t
  25     | Leaf
  26
  27   let empty = Leaf
  28
  29   let rec set t ~k ~v =
  30     match t with
  31     | Leaf -> Node (k, v, Leaf, Leaf)
  32     | Node (k', v', l, r) when k < k' -> Node (k', v', set l ~k ~v, r)
  33     | Node (k', v', l, r) when k > k' -> Node (k', v', l, set r ~k ~v)
  34     | Node (k, _, l, r) -> Node (k, v, l, r)
  35
  36   let rec get t ~k =
  37     match t with
  38     | Leaf -> None
  39     | Node (k', _, l, _) when k < k' -> get l ~k
  40     | Node (k', _, _, r) when k > k' -> get r ~k
  41     | Node (_, v, _, _) -> Some v
  42
  43   let rec member t ~k =
  44     match t with
  45     | Leaf -> false
  46     | Node (k', _, l, _) when k < k' -> member l ~k
  47     | Node (k', _, _, r) when k > k' -> member r ~k
  48     | Node _ -> true
  49
  50   let rec fold t ~f ~init:acc =
  51     match t with
  52     | Leaf -> acc
  53     | Node (k, v, l, r) -> fold r ~f ~init:(f (fold l ~f ~init:acc) (k, v))
  54
  55   let print t ~k_to_string ~indent_level =
  56     let (_, nodes) =
  57       fold t
  58         ~f:(fun (indentation, nodes) (k, _) ->
  59             let indentation = indent_level ^ indentation in
  60             let node = indentation ^ (k_to_string k) in
  61             (indentation, node :: nodes)
  62           )
  63         ~init:(indent_level, [])
  64     in
  65     List.iter (List.rev nodes) ~f:print_endline
  66 end
  67
  68 let () =
  69   let tree_a = BinaryTree.empty in
  70   let tree_a = BinaryTree.set tree_a ~k:"k1" ~v:"v1" in
  71   let tree_a = BinaryTree.set tree_a ~k:"k2" ~v:"v2" in
  72   assert (BinaryTree.member tree_a ~k:"k1");
  73   assert (BinaryTree.member tree_a ~k:"k2");
  74   assert (Some "v1" = BinaryTree.get tree_a ~k:"k1");
  75   assert (Some "v2" = BinaryTree.get tree_a ~k:"k2");
  76   let tree_b =
  77     Array.fold_left
  78       (Sys.argv)
  79       ~init:BinaryTree.empty
  80       ~f:(fun t k -> BinaryTree.set t ~k ~v:())
  81   in
  82   Printf.printf "%B\n" (BinaryTree.member tree_b ~k:"a");
  83   BinaryTree.print tree_b ~k_to_string:(fun x -> x) ~indent_level:"-";