Simple, Efficient Trie Maps

I became interested in implementing Bentley and Sedgewick's Ternary Trees by reading the challenge proposed on Programming Praxis (a great blog, by the way). I thought the algorithms would be a good fit for a purely functional data structure, and I am happy with the results. Here is my implementation.

I follow the Map.S signature, so that this can be a drop-in replacement for string-keyed maps, so that the interface (.mli file) is simply:

include Map.S with type key = string

For the implementation I restrict the key space to strings:

type key = string

Trees are optional references to nodes, and these are records containing the mapped value whenever this node represents a full key, the character suffix for the tree and the three branches:

type 'a t = 'a node option
and 'a node =
{ value : 'a option;
  split : char;
  lokid : 'a t;
  eqkid : 'a t;
  hikid : 'a t; }

In this I follow closely the paper, except for a crucial detail: in it, the authors make use of the terminator inherent to every C string, the null byte '\0'. This forces a design decision:

• emulate the termination character. This has the distinct disadvantage of having to prohibit it from occurring in string keys
• use a char option as the key. This forces boxing the key into a constructed type and complicates the logic in all the comparisons, which has a negative impact in the inner loops
• use a char list or string as the key. This has the advantage that it can be adapted to compress the search paths, but it requires generating substrings of everything
• guard on the string length. This is the choice I took. It complicates the comparison logic somewhat, but it has the advantage of making the node representation more regular and compact, and allowing the full range of characters in string keys

An empty tree is represented by a None reference:

let empty = None

Consequently, the test for an empty tree is a match on the option:

let is_empty = function None -> true | Some _ -> false

This means that I'll be forced to normalize trees on deletion, of which more later. The workhorse function is lookup:

let lookup k t =
  let n = String.length k in
  let rec go i = function
  | None   -> None
  | Some t ->
    let cmp = Char.compare k.[i] t.split in
    if cmp < 0 then go  i    t.lokid else
    if cmp > 0 then go  i    t.hikid else
    if i+1 < n then go (i+1) t.eqkid else
                    t.value
  in go 0 t

It is as simple as it can be given the choices I outlined above. The crucial bit is that the presence or absence of the value is the flag that indicates if this node represents a present or absent key in the tree. The signature's retrieval and membership functions are simply wrappers around this:

let find k t = match lookup k t with
| Some x -> x
| None   -> raise Not_found

let mem k t = match lookup k t with
| Some _ -> true
| None   -> false

Inserting a key-value pair in the tree is a bit more involved:

let add k v t =
  let n = String.length k in
  let rec go i = function
  | None when i+1 < n ->
    Some { value = None  ; split = k.[i]; lokid = None; eqkid = go (i+1) None; hikid = None }
  | None              ->
    Some { value = Some v; split = k.[i]; lokid = None; eqkid =          None; hikid = None }
  | Some t            ->
    let cmp = Char.compare k.[i] t.split in
    if cmp < 0 then Some { t with lokid = go  i    t.lokid } else
    if cmp > 0 then Some { t with hikid = go  i    t.hikid } else
    if i+1 < n then Some { t with eqkid = go (i+1) t.eqkid } else
                    Some { t with value = Some v }
  in go 0 t

The case when the end of a branch is reached is inlined, so that a descending branch keyed on the key's suffix is built until reaching its last character, which stores the keyed value. Otherwise, the suffix is compared to the node key, and the appropriate child is built. Note that the recursively built subtree replaces the corresponding child in the node; this is what makes this data structure purely functional and persistent. Removing a mapping is similar, but complicated by the need to normalize the tree so that an empty tree is represented by None:

let remove k t =
  let prune = function
  | { value = None; lokid = None; eqkid = None; hikid = None } -> None
  | t -> Some t
  in
  let n = String.length k in
  let rec go i = function
  | None   -> None
  | Some t ->
    let cmp = Char.compare k.[i] t.split in
    if cmp < 0 then prune { t with lokid = go  i    t.lokid } else
    if cmp > 0 then prune { t with hikid = go  i    t.hikid } else
    if i+1 < n then prune { t with eqkid = go (i+1) t.eqkid } else
                    prune { t with value = None             }
  in go 0 t

In this case, prune works somewhat as a smart constructor for node references.

The big advantage of ternary search trees is that, since they are modeled on three-way quicksort, they provide ordered key storage with relative insensitivity to the order of insertions. The big disadvantage is that they don't have a canonical form for a given key set, that is, their shape depends on the order in which the keys are inserted. This means that they provide a cheap fold but no easy way to compare trees:

let fold f t e =
  let rec go prefix e = function
  | None   -> e
  | Some t ->
    let e   = go prefix e t.lokid in
    let key = prefix ^ String.make 1 t.split in
    let e   = match t.value with None -> e | Some v -> f key v e in
    let e   = go key    e t.eqkid in
              go prefix e t.hikid
  in go "" e t

Note that the key is built character by character but only on the equal kids. Note also that the resulting fold is guaranteed to be ordered on the keys by the sequence in which the recursion is carried out. The rest of the iteration functions in the signature are trivial wrappers:

let iter f t = fold (fun k v () ->        f k v ) t ()
and map  f t = fold (fun k v    -> add k (f v  )) t empty
and mapi f t = fold (fun k v    -> add k (f k v)) t empty

As I indicated above, ternary search trees have no canonical form, so to compare two trees you have to somehow traverse them in synchronization. One easy but expensive way would be to turn both into key-value lists and compare those; this would require O(n min m) time and O(n + m) space. The best way would be to turn the fold into a continuation-passing iterator that lazily streams keys and traverse both in tandem. I leave it as an exercise to the reader:

let compare _ _ _ = failwith "Not comparable"

let equal eq t t' =
  0 == compare (fun v w -> if eq v w then 0 else 1) t t'