#light

type Colour = | Red | Black
type 'a Tree =
   | Node of 'a * 'a Tree * 'a Tree * Colour
   | Leaf

let rec balance = function
   | Node(y, Node(x, a, b, Red), Node(z, c, d, Red), Black)
   | Node(x, a, Node(z, Node(y, b, c, Red), d, Red), Black)
   | Node(x, a, Node(y, b, Node(z, c, d, Red), Red), Black)
   | Node(z, Node(x, a, Node(y, b, c, Red), Red), d, Black)
   | Node(z, Node(y, Node(x, a, b, Red), c, Red), d, Black) -> Node(y, Node(x, a, b, Black), Node(z, c, d, Black), Red)
   | a -> a

let makeblack = function
   | Node(i, l, r, _) -> Node(i, l, r, Black)
   | Leaf -> Leaf

let makered (Node(i, l, r, _)) = Node(i, l, r, Red)


let rec search tree value =
   match tree with
   | Leaf -> None
   | Node(i, l, _, _) when value < i -> search l value
   | Node(i, _, r, _) when value > i -> search r value
   | _ -> Some value


let add tree value =
   let rec insert = function
   | Leaf -> Node(value, Leaf, Leaf, Red)
   | Node(i, l, r, c) when value < i -> Node(i, insert l, r, c) |> balance
   | Node(i, l, r, c) when value > i -> Node(i, l, insert r, c) |> balance
   | a -> a

   insert tree |> makeblack


let remove tree value =
   let number = ref None
   let rec rem = function

      //red -> (black), black
      | Node(i, Node(j, jl, jr, Black), Node(k, kl, kr, Black), Red) when value < i ->
         Node(i, Node(j, jl, jr, Red) |> rem, Node(k, kl, kr, Red), Black) |> balance
      | Node(i, Node(j, jl, jr, Black), Node(k, kl, kr, Black), Red) when value > i ->
         Node(i, Node(j, jl, jr, Red), Node(k, kl, kr, Red) |> rem, Black) |> balance

      //red -> black/red/leaf, (red):
      | Node(p, l, Node(c, cl, cr, Red), Red) when value > p -> Node(p, l, Node(c, cl, cr, Red) |> rem, Red)
      | Node(p, Node(c, cl, cr, Red), r, Red) when value < p -> Node(p, Node(c, cl, cr, Red) |> rem, r, Red)

      //red -> (black), red
      | Node(n3, Node(n1, Node(n0, l0, r0, Black), Node(n2, l2, r2, Black), Red), n4, Red) when value > n3 ->
         Node(n1, Node(n0, l0, r0, Black), Node(n3, Node(n2, l2, r2, Black), n4, Red) |> rem, Red)
      | Node(n1, n0, Node(n3, Node(n2, l2, r2, Black), Node(n4, l4, r4, Black), Red), Red) when value < n1 ->
         Node(n3, Node(n1, n0, Node(n2, l2, r2, Black), Red) |> rem, Node(n4, l4, r4, Black), Red)

      //(red) -> leaf
      | Node(i, Leaf, Leaf, Red) when value = i -> Leaf
      | Node(i, l, Leaf, Red) when value = i -> l

      //non-leaf match:
      | Node(i, l, r, Red) as a when value = i ->
         //find successor:
         match a with
         | Node(i, Node(j, jl, jr, Black), Node(k, kl, kr, Black), Red) ->
            Node(i, Node(j, jl, jr, Red), Node(k, kl, kr, Red) |> remsuccessor, Black) |> balance
         | Node(p, l, Node(c, cl, cr, Red), Red) -> Node(p, l, Node(c, cl, cr, Red) |> remsuccessor, Red)
         | Node(n3, Node(n1, Node(n0, l0, r0, Black), Node(n2, l2, r2, Black), Red), n4, Red) ->
           Node(n1, Node(n0, l0, r0, Black), Node(n3, Node(n2, l2, r2, Black), n4, Red) |> rem, Red)
   and remsuccessor = function
      | Node(i, Node(j, jl, jr, Black), Node(k, kl, kr, Black), Red) ->
         Node(i, Node(j, jl, jr, Red) |> remsuccessor, Node(k, kl, kr, Red), Black) |> balance
      | Node(p, Node(c, cl, cr, Red), r, Red) -> Node(p, Node(c, cl, cr, Red) |> remsuccessor, r, Red)
      | Node(n1, n0, Node(n3, Node(n2, l2, r2, Black), Node(n4, l4, r4, Black), Red), Red) ->
         Node(n3, Node(n1, n0, Node(n2, l2, r2, Black), Red) |> remsuccessor, Node(n4, l4, r4, Black), Red)
      | Node(i, Leaf, Leaf, Red) -> number := Some i; Leaf
      | Node(i, Leaf, r, Red) -> number := Some i; r
   and check tree =
      match !number with
      | Some x -> switch x tree
      | None -> tree
   and switch x = function
      | Node(i, l, r, c) when value < i -> Node(i, switch x l, r, c)
      | Node(i, l, r, c) when value > i -> Node(i, l, switch x r, c)
      | Node(i, l, r, c) when value = i -> Node(x, l, r, c)

   tree |> makered |> rem |> makeblack |> check


let rec traverse tree = seq {
   match tree with
   | Leaf -> ()
   | Node(i, l, r, _) ->
      yield! traverse l
      yield i
      yield! traverse r
}


let sort x = x |> Seq.fold add Leaf |> traverse
 

previous: binary search tree