trie.ml

(*
Let us define a trie using two mutually defined types (given in the prelude):

* trie which represents a trie, that is a tree whose root may contain an integer and whose children are indexed by characters ;
* char_to_children which implements the associative data structure whose keys are characters and whose values are trie (childrens).
*
*)

type trie = Trie of int option * char_to_children
and char_to_children = (char * trie) list;;


(*
As a trade-off between speed and memory consumption, we choose an associative
list to represent the association between characters and children. 

The prelude also gives examples of empty trie and of another one that contains
the following pairs (key, value): 

[("A", 15); ("to", 7); ("tea", 3);("ted", 4); ("ten", 12); ("i", 11); ("in", 5); ("inn", 9)].

*)

let empty =
    Trie (None, []);;

let example =
  Trie (None,
	[('i', Trie (Some 11,
                     [('n', Trie (Some 5, [('n', Trie (Some 9, []))]))]));
	 ('t',
	  Trie (None,
		[('e',
		  Trie (None,
			[('n', Trie (Some 12, [])); ('d', Trie (Some 4, []));
			 ('a', Trie (Some 3, []))]));
		 ('o', Trie (Some 7, []))]));
     ('A', Trie (Some 15, []))]);;


(*
Write a function children_from_char : char_to_children -> char -> trie option such that
    children_from_char m c = Some t if (c, t) is the first pair in m with c as a first component ;
    children_from_char m c = None if no such pair exists in m.
*)

let children_from_char (c2c : char_to_children) (c : char) : (trie option) =
    try Some (List.assoc c c2c) with
    | exn -> None
;;

let get_single_layer c (Trie (i, childrans)) = children_from_char childrans c;;

get_single_layer 'A'  example;;


(*
Write a function `update_children : char_to_children -> char -> trie -> char_to_children` such that
    1. `children_from_char (update_children m c t) c = Some t` ;
    2. `children_from_char (update_children m c t) c' = children_from_char m c'`
        for `c <> c'`;
    3. If `children_from_char m c = Some t` then
       `List.length (update_children m c t') = List.length m`.
*)

let rec add_association key value assoc_list = match assoc_list with
| [] -> [(key, value)]
| (k, v)::rst when k = key -> (key, value)::rst
| (k, v)::rst when k != key -> (k, v)::(add_association key value rst)
| _ -> []
;;


let update_children (children : char_to_children) (c : char) (tt : trie) : char_to_children =
    add_association c tt children;;


(*
Write a function lookup : trie -> string -> int option such that lookup trie w = Some i if i is the value of the key w in trie and lookup trie w = None if w is not a key of trie. 

To look for a key in a trie, iterate over the characters of the key from left to right. Given the current character c and the current node of the trie n, look for the children n for character c. If such a children exists, continue with that trie and the remainder of the key. If no such children exists, the key is not in the trie. When the characters of the key are entirely consumed, look at the root of the current trie. If there is an integer, this is the value you are looking for. If there is no integer, the key not in the trie.

*)

let split s = match String.length s with
| 0 -> None
| n -> Some (String.get s 0, String.sub s 1 (n - 1))
;;


let rec lookup (Trie (v, childrens)) (key : string) : int option =  match split key with
| None -> v
| Some (hd, rst) -> match children_from_char childrens hd with
                    | None -> None
                    | Some next_trie -> lookup next_trie rst
;;

(*
Write a function insert : trie -> string -> int -> trie such that lookup (insert trie w k) w = Some k and lookup (insert trie w k) w' = lookup trie w' for w <> w'.
*)

let rec insert (Trie (old_v, childrens)) (key : string) (v : int) : trie = match split key with
| None -> Trie (Some v,  childrens)
| Some (k, rst) -> (match children_from_char childrens k with
                   | None -> Trie (old_v, update_children childrens k (insert empty rst v))
                   | Some next_trie -> Trie (old_v, update_children childrens k (insert next_trie rst v))
                   )
;;