Trie

trie.ml

(* Simple trie to store strings by Nick Heiner <[email protected]> *)

module CharMap = Map.Make(Char)

(* count of members of the set that end at this node * mapping from
   next char => children *)
type trie = Node of int * trie CharMap.t

let empty = Node (0, CharMap.empty)

(* Add a string to the trie *)
let rec add (Node (count, children) : trie) = function 
  | "" -> Node (count + 1, children)
  | str -> 
    let firstChar = String.get str 0 in
    let lastChars = String.sub str 1 ((String.length str) - 1) in
    let newTrie = if CharMap.mem firstChar children
      then add (CharMap.find firstChar children) lastChars
      else add empty lastChars in
    let newChildren = CharMap.add firstChar newTrie children in   
    Node (count, newChildren)
    
let addAll (trie : trie) (strs : string list) : trie =
  List.fold_left (fun trieAcc str -> add trieAcc str) empty strs 
  
(* Get all the strings of a trie *)
let traverse (trie : trie) : string list =
  let rec helper (Node (count, children) : trie) (prevChar : char option)
      : string list =
    let string_of_char chr = 
      match chr with 
        | None -> ""
        | Some ch -> String.make 1 ch in
    let perChild = List.map (fun (chr, child) -> 
      (helper child (Some chr))) (CharMap.bindings children) in
          let fromChildren = List.flatten perChild in
    let withPrev =
      List.map (fun str -> (string_of_char prevChar) ^ str) fromChildren in
    let rec clone str count = 
      if count = 0 then [] else str::(clone str (count - 1))
    in (clone (string_of_char prevChar) count) @ withPrev in  
  helper trie None

let sort (strs: string list): string list =
  traverse (addAll empty strs)