[LeetCode] 208. Implement Trie (Prefix Tree)

Implement a trie with insert, search, and startsWith methods.

Example:

1 2 3 4 5 6 7 8

Trie trie = new Trie(); trie.insert("apple"); trie.search("apple"); // returns true trie.search("app"); // returns false trie.startsWith("app"); // returns true trie.insert("app"); trie.search("app"); // returns true

Note:

• You may assume that all inputs are consist of lowercase letters a-z.
• All inputs are guaranteed to be non-empty strings.

Solution

A Trie Tree(Prefix Tree) is mainly used to retrieve a key from the dictionary. Things like Autocomplete in search engine and spell checker utilize this data structure extensively. A more detailed explanation can be found here.

In a Trie Tree, each node has its links and a label indicating the end of a word. In my implementation, I used a map structure to store the links. Note that this implementation affects the performance of the operations on the Tree since the map is actually a linked-list of tuples instead of traditional arrays. Haskell do have array-like types (Data.Array, Data.Seq…), but here let’s focus more on the concept of a Trie Tree.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

data TrieNode = Node [(Char, TrieNode)] Bool deriving (Show, Read)

root = Node [] False

tip = Node [] True

insert :: String -> TrieNode -> TrieNode
insert [] (Node links end) = tip
insert (a : str) (Node links end)
  | null $ fst match = Node ((a, insert str root) : links) end
  | otherwise = Node ((a, insert str $ snd . head . fst $ match) : snd match) end
  where match = partition (\x -> a == fst x) links

search :: String -> TrieNode -> Bool
search []        (Node _     end) = end
search (a : str) (Node links _  ) = case lookup a links of
  Nothing -> False
  Just x  -> search str x

startsWith :: String -> TrieNode -> Bool
startsWith []        _              = True
startsWith (a : str) (Node links _) = case lookup a links of
  Nothing -> False
  Just x  -> startsWith str x

In the insert function, basically what we do is to check the links of a node. If there is already a link for the character we would like to insert into, then we need to insert the rest of our string into the node that particular link points to. Since everything in Haskell is immutable, we must reconstruct a list of links for the node. On the other hand, if we are going to insert a character that is not presented in the links, we just prepend a new link to the list.

Searching and finding prefix is pretty straightforward, we search the character in the links, if there is a hit then we proceed to the node it’s pointing to and search for the next one and so on. The difference is when checking a prefix, we don’t care if the node marks the end of a word.

Complexity Analysis

For insertion, the time complexity is \(O(m)\), where \(m\) is the length of the word. The space complexity is \(O(m)\).

For searching, it is obvious that the time complexity is \(O(m)\) and the space complexity is \(O(1)\).