Commits

Luke Plant  committed 862ca93

Use a Map instead of list to store children in word tree

This results in simpler and faster code.

  • Participants
  • Parent commits 81d6e69

Comments (0)

Files changed (1)

File phone2word.hs

     then Just char
     else Map.lookup (toUpper char) numbersMap
 
-data WordTree = Node { nodeLabel :: Char           -- digit
-                     , nodeSubForest :: [WordTree] -- children
-                     , nodeWords :: [B.ByteString] -- words that correspond to this node
+data WordTree = Node { nodeSubForest :: Map.Map Char WordTree -- children
+                     , nodeWords :: [B.ByteString]            -- words that correspond to this node
                      } deriving (Eq, Read, Show)
 
 
- -- Look up a char in the sub forests, and return ([matching WordTree], [non-matching WordTree])
-subForestLookup :: Char -> WordTree -> ([WordTree],[WordTree])
-subForestLookup char tree = partition (\n -> nodeLabel n == char) $ nodeSubForest tree
-
 --
 -- Building the tree
 --
 
-startNode = Node { nodeLabel = 'X' -- never used
-                 , nodeSubForest = []
+startNode = Node { nodeSubForest = Map.empty
                  , nodeWords = []
                  }
 
 newNode :: Char -> WordTree
-newNode digit = Node { nodeLabel = digit
-                     , nodeSubForest = []
+newNode digit = Node { nodeSubForest = Map.empty
                      , nodeWords = []
                      }
 
             in case md of
                  Nothing -> tree -- no corresponding digit to letter, can't add word
                  Just d ->
-                     let (matching, others) = subForestLookup d tree
+                     let matching = Map.lookup d $ nodeSubForest tree
                          node = case matching of
                                   -- if node does not exist, need to create it
-                                  [] -> newNode d
-                                  [n] -> n
-                                  _ -> error "impossible happened"
+                                  Nothing -> newNode d
+                                  Just n -> n
                          node' = addWord cs word node
-                     in tree { nodeSubForest = node' : others }
+                     in tree { nodeSubForest = Map.insert d node' $ nodeSubForest tree }
 
 addWordToList w ws =
     -- eliminate duplicates
        then nodeWords tree
        else let d = B.head chars
                 ds = B.tail chars
-                (ns, others) = subForestLookup d tree
-            in case ns of
-                 [] -> []
-                 (n:_) -> findWords ds n
+                nd = Map.lookup d $ nodeSubForest tree
+            in case nd of
+                 Nothing -> []
+                 Just n -> findWords ds n
 
 --
 -- Search for combos
 getWords :: String -> IO [B.ByteString]
 getWords dictFile = B.readFile dictFile >>= (return . B.lines)
 
-countNodes t = 1 + (sum $ map countNodes (nodeSubForest t))
+countNodes t = 1 + (sum $ map countNodes (Map.elems $ nodeSubForest t))
 
-countWords t = length (nodeWords t) + (sum $ map countWords (nodeSubForest t))
+countWords t = length (nodeWords t) + (sum $ map countWords (Map.elems $ nodeSubForest t))
 
 main = do
   dictFile:_ <- getArgs