Map implementation in F#

Question

Just as a refresher I put together a simple Map implementation and I would love to get some feedback on it.

open System
open System.Collections.Generic
type Node<'a, 'b when 'a : comparison> = {
 key : 'a
 value : 'b
 left : option<Node<'a, 'b>> 
 right : option<Node<'a, 'b>>
}
type Map<'a, 'b when 'a : comparison>(root : option<Node<'a, 'b>>) =
 let comparer = LanguagePrimitives.FastGenericComparer<'a>
 let rec add (key : 'a) (value : 'b) (node : Node<'a, 'b>) =
 match comparer.Compare (key, node.key) with
 | r when r < 0 ->
 match node.left with
 | Some n ->
 match comparer.Compare (key, n.key) with
 | r when r > 0 ->
 let left = Some { 
 key = key
 value = value 
 left = node.left 
 right = None 
 }
 Some { node with left = left }
 | _ ->
 Some { node with left = add key value n }
 | None ->
 let left = Some { 
 key = key
 value = value 
 left = None 
 right = None 
 }
 Some { node with left = left }
 | r when r > 0 ->
 match node.right with
 | Some n ->
 match comparer.Compare (key, n.key) with
 | r when r < 0 ->
 let right = Some { 
 key = key
 value = value 
 left = None 
 right = node.right 
 }
 Some { node with right = right }
 | _ ->
 Some { node with right = add key value n }
 | None ->
 let right = Some { 
 key = key
 value = value 
 left = None 
 right = None 
 }
 Some { node with right = right } 
 | _ -> 
 Some { node with value = value }
 let rec find (key : 'a) (node : Node<'a, 'b>) =
 match comparer.Compare (key, node.key) with
 | r when r < 0 ->
 match node.left with
 | Some node -> find key node
 | None -> raise (KeyNotFoundException())
 | r when r > 0 ->
 match node.right with
 | Some node -> find key node
 | None -> raise (KeyNotFoundException())
 | _ -> node.value 
 member x.Item key =
 match root with
 | Some node -> 
 find key node
 | None -> raise (KeyNotFoundException()) 
 member x.Add (key : 'a, value : 'b) =
 match root with
 | Some node -> 
 Map (add key value node)
 | None ->
 let node = Some { 
 key = key 
 value = value 
 left = None
 right = None 
 } 
 Map node
 static member Empty = Map<'a, 'b>(None)
 static member FromSeq (s:seq<'a * 'b>) =
 s 
 |> Seq.fold 
 (fun (m:Map<'a, 'b>) (k, v) -> m.Add (k, v)) 
 Map<'a, 'b>.Empty

After running some basic tests I found my implementation is comparable to the FSharpMapclass in performance and sometimes better. Obviously I don't have the time/desire to test this extensively so take that with a grain of salt. I'm wondering if anyone can spot a characteristic of this code that will cause a performance breakdown under certain conditions or for certain types of keys.

Improved Version

open System
open System.Collections.Generic
type Node<'a, 'b when 'a : comparison> = {
 key : 'a
 value : 'b
 height : int
 left : option<Node<'a, 'b>> 
 right : option<Node<'a, 'b>>
}
type Map<'a, 'b when 'a : comparison>(root : option<Node<'a, 'b>>) =
 let comparer = LanguagePrimitives.FastGenericComparer<'a> 
 let height node =
 match node with
 | Some node -> node.height
 | None -> 0 
 let make key value left right =
 let h = 
 match height left, height right with
 | l, r when l >= r -> l + 1
 | l, r -> r + 1
 Some { 
 key = key; 
 value = value; 
 height = h; 
 left = left; 
 right = right 
 } 
 let balance key value left right =
 match height left, height right with
 | l, r when r > l + 2 ->
 match right with
 | Some rn ->
 match height rn.left with
 | rl when rl <= l + 1 ->
 let left = make key value left rn.left 
 make rn.key rn.value left rn.right
 | _ ->
 match rn.left with
 | Some rnl ->
 let left = make key value left rnl.left
 let right = make rn.key rn.value rnl.right rn.right
 make rnl.key rnl.value left right
 | None ->
 make key value left right
 | None -> 
 make key value left right
 | l, r when l <= r + 2 -> 
 make key value left right
 | l, r ->
 match left with
 | Some ln ->
 match height ln.right with
 | rl when rl <= l + 1 ->
 let right = make key value ln.right right 
 make ln.key ln.value ln.left right
 | _ ->
 match ln.right with
 | Some lnr ->
 let left = make ln.key ln.value ln.left lnr.left
 let right = make key value lnr.right right
 make lnr.key lnr.value left right
 | None ->
 make key value left right
 | None -> 
 make key value left right
 let rec add key value node =
 match comparer.Compare (key, node.key) with
 | r when r < 0 ->
 match node.left with
 | Some n -> 
 balance node.key node.value (add key value n) node.right
 | None ->
 let left = Some { 
 key = key
 value = value
 height = node.height + 1
 left = None
 right = None 
 }
 balance node.key node.value left node.right
 | r when r > 0 ->
 match node.right with
 | Some n ->
 balance node.key node.value node.left (add key value n)
 | None ->
 let right = Some { 
 key = key
 value = value
 height = node.height + 1
 left = None
 right = None 
 }
 balance node.key node.value node.left right
 | _ -> 
 Some { node with value = value }
 let rec find key node =
 match comparer.Compare (key, node.key) with
 | r when r < 0 ->
 match node.left with
 | Some node -> find key node
 | None -> raise (KeyNotFoundException())
 | r when r > 0 ->
 match node.right with
 | Some node -> find key node
 | None -> raise (KeyNotFoundException())
 | _ -> node.value 
 member x.Item key =
 match root with
 | Some node -> 
 find key node
 | None -> raise (KeyNotFoundException()) 
 member x.Add (key, value) =
 match root with
 | Some node -> 
 Map (add key value node)
 | None ->
 let node = Some { 
 key = key 
 value = value 
 height = 1
 left = None
 right = None 
 } 
 Map node
 static member Empty = Map<'a, 'b>(None)
 static member OfSeq (s:seq<'a * 'b>) =
 s 
 |> Seq.fold 
 (fun (m:Map<'a, 'b>) (k, v) -> m.Add (k, v)) 
 Map<'a, 'b>.Empty

Answer 1

Second revision

Ok, first of all it would really help if you added some comments. I would also suggest giving your variables names which have more than one letter.

In particular it would be a good idea to describe the algorithms being used. In particular you should document what kind of tree and balancing algorithm you're using.

Speaking of: your tree kind of looks like an AVL tree except that in AVL trees the maximum difference between heights is 1, while you allow a difference of 2. Since you never said that you were trying to implement an AVL tree, I'm not sure whether that's intentional or not.

---

Some notes on particular pieces of code:

match height left, height right with
| l, r when l >= r -> l + 1
| l, r -> r + 1

That's just a cumbersome way to write max (height left) (height right).

---

Some {
 key = key
 value = value
 height = node.height + 1
 left = None
 right = None
}

This looks like a mistake. If both the subtrees are empty, then clearly the height is 1 and not dependent on node.

---

I would also recommend that instead of using Node options you define an actual tree type like this: type ('a * 'b) tree = Node of ('a, 'b) Node | EmptyTree. This way you can refer to empty trees as EmptyTree rather than None and to nodes as Node {...} rather than Some {...}. Of course that's just cosmetics, but I do think it reads much nicer.

---

First revision

One minor style point is that you might want to call your FromSeq method ofSeq instead as that is what the equivalent function of F#'s Map module is called.

---

Regarding performance the most obvious problem is that your tree is not balanced in any way. If you create a map from a sorted list of keys, it will degenerate into a linked list and have much worse performance than F#'s standard map class. Just compare the time it takes to create your map from Seq.zip [1..500000] [1..500000] to the time the standard Map needs for the same input.

To address this, you should implement some kind of balancing. For example you may use red-black trees, which is what F#'s standard map uses.

• ChaosPandion: You were so right about my implementation. I tried adding 50k keys and it exploded with a StackOverflowException. After I implemented balancing, my implementation comes out to be oh so slightly faster the the built-in version which I crop up to slightly less object creation and less robust code. Anyway, thanks so much and feel free to comment on my latest version.
• sepp2k: @Chaos: I've edited my answer in response to your revision.
• ChaosPandion: Damn, you're right. I can't half ass this. I'm gonna tidy things up and drop you a comment sometime next week just in case you are curious.