public
Last active

naive trie map in haskell

  • Download Gist
trie.hs
Haskell
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
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

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.