Skip to content

Instantly share code, notes, and snippets.

@yihuang
Created February 27, 2012 12:46
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save yihuang/1923486 to your computer and use it in GitHub Desktop.
naive trie map in haskell
import Prelude hiding (lookup)
import Data.Maybe (fromMaybe)
import qualified Data.Map as M
import Control.Monad (join)
data Trie c a = Trie
(Maybe a) -- node may or may not contain a value.
(M.Map c (Trie c a))
deriving (Show)
empty :: Trie c a
empty = Trie Nothing M.empty
insert :: Ord c => [c] -> a -> Trie c a -> Trie c a
insert [] a (Trie _ m) = Trie (Just a) m
insert (c:cs) a (Trie v m) = Trie v m'
where
m' = M.alter alt c m
alt = Just . insert cs a . fromMaybe empty
fromList :: Ord c => [([c], a)] -> Trie c a
fromList = foldr (uncurry insert) empty
addPrefix :: c -> [([c], a)] -> [([c], a)]
addPrefix c = map (mapFst (c:))
where mapFst f (a, b) = (f a, b)
toList :: Trie c a -> [([c], a)]
toList (Trie mv m) =
case mv of
Nothing -> rest
Just v -> ([], v) : rest
where
rest = flip concatMap (M.assocs m) $
\(c, t) -> addPrefix c (toList t)
prefix :: Ord c => [c] -> Trie c a -> [([c], a)]
prefix [] t = toList t
prefix (c:cs) (Trie _ m) =
maybe [] (addPrefix c . prefix cs) $ M.lookup c m
allPrefixes :: Ord c => Trie c a -> [[c]]
allPrefixes = map fst . toList
lookup :: Ord c => [c] -> Trie c a -> Maybe a
lookup [] (Trie mv _) = mv
lookup (c:cs) (Trie _ m) = join $ fmap (lookup cs) $ M.lookup c m
main = do
let t = toList $ fromList $
[ ("how", 1)
, ("when", 2)
, ("what", 3)
, ("where", 4)
]
print t
-- toList and fromList is revertable, which can be auto-checked by QuickCheck.
print $ toList $ fromList t
print $ prefix "w" $ fromList t
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment