#undef

1 files changed, 40 insertions, 0 deletions

diff --git a/tree.sml b/tree.sml
index 57e9a24..0a788af 100644
--- a/tree.sml
+++ b/tree.sml
@@ -1,6 +1,8 @@
 structure Tree: TREE = struct
   datatype ('k, 'v) t = Node of 'k * 'v * ('k, 'v) t * ('k, 'v) t | Empty
 
+  type 'k cmp = 'k -> 'k -> order
+
   val empty = Empty
 
   fun insert _ Empty k v = (NONE, Node (k, v, Empty, Empty))
@@ -20,6 +22,44 @@ structure Tree: TREE = struct
             (res, Node (k', v', left, right))
           end
 
+  fun delete _ Empty _ = (NONE, Empty)
+    | delete cmp (Node (k', v', left, right)) k =
+        case cmp k k' of
+          LESS =>
+          let
+            val (res, left) = delete cmp left k
+          in
+            (res, Node (k', v', left, right))
+          end
+        | GREATER =>
+          let
+            val (res, right) = delete cmp right k
+          in
+            (res, Node (k', v', left, right))
+          end
+        | EQUAL => (
+            case (left, right) of
+              (Empty, Node _) => (SOME v', right)
+            | (Node _, Empty) => (SOME v', left)
+            | (Empty, Empty) => (SOME v', Empty)
+            | (Node _, Node _) =>
+              let
+                fun deleteRightmost Empty = raise Unreachable
+                  | deleteRightmost (Node (k, v, left, Empty)) =
+                    ((k, v), left)
+                  | deleteRightmost (Node (k, v, left, right)) =
+                  let
+                    val (p, right) = deleteRightmost right
+                  in
+                    (p, Node (k, v, left, right))
+                  end
+
+                val ((k, v), left) = deleteRightmost left
+              in
+                (SOME v', Node (k, v, left, right))
+              end
+        )
+
   fun lookup _ Empty _ = NONE
     | lookup cmp (Node (k', v', left, right)) k =
         case cmp k k' of