naive trie map in haskell
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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