exercises/ch01/tree.ml

   1 module Array = ArrayLabels
   2
   3 module type TREE = sig
   4   type ('k, 'v) t
   5
   6   val empty : ('k, 'v) t
   7
   8   val set : ('k, 'v) t -> k:'k -> v:'v -> ('k, 'v) t
   9
  10   val get : ('k, 'v) t -> k:'k -> 'v option
  11
  12   val member : ('k, 'v) t -> k:'k -> bool
  13 end
  14
  15 module BinaryTree : TREE = struct
  16   type ('k, 'v) t =
  17     | Node of 'k * 'v * ('k, 'v) t * ('k, 'v) t
  18     | Leaf
  19
  20   let empty = Leaf
  21
  22   let rec set t ~k ~v =
  23     match t with
  24     | Leaf -> Node (k, v, Leaf, Leaf)
  25     | Node (k', v', l, r) when k < k' -> Node (k', v', set l ~k ~v, r)
  26     | Node (k', v', l, r) when k > k' -> Node (k', v', l, set r ~k ~v)
  27     | Node (k, _, l, r) -> Node (k, v, l, r)
  28
  29   let rec get t ~k =
  30     match t with
  31     | Leaf -> None
  32     | Node (k', _, l, _) when k < k' -> get l ~k
  33     | Node (k', _, _, r) when k > k' -> get r ~k
  34     | Node (_, v, _, _) -> Some v
  35
  36   let rec member t ~k =
  37     match t with
  38     | Leaf -> false
  39     | Node (k', _, l, _) when k < k' -> member l ~k
  40     | Node (k', _, _, r) when k > k' -> member r ~k
  41     | Node _ -> true
  42 end
  43
  44 let () =
  45   let tree_a = BinaryTree.empty in
  46   let tree_a = BinaryTree.set tree_a ~k:"k1" ~v:"v1" in
  47   let tree_a = BinaryTree.set tree_a ~k:"k2" ~v:"v2" in
  48   assert (BinaryTree.member tree_a ~k:"k1");
  49   assert (BinaryTree.member tree_a ~k:"k2");
  50   assert (Some "v1" = BinaryTree.get tree_a ~k:"k1");
  51   assert (Some "v2" = BinaryTree.get tree_a ~k:"k2");
  52   let tree_b =
  53     Array.fold_left
  54       (Sys.argv)
  55       ~init:BinaryTree.empty
  56       ~f:(fun t k -> BinaryTree.set t ~k ~v:())
  57   in
  58   Printf.printf "%B\n" (BinaryTree.member tree_b ~k:"a")