haskell_generic_trie.hs

module KTrie ( contains 
             , empty 
             , fullDeepLookup 
             , insert 
             , lookup 
             , null 
             , remove 
             , toList 
             , Trie ) where 
 
import qualified Data.Map as M 
import Data.Map ((!)) 
import Prelude hiding (null, lookup) 
import Data.Maybe (fromMaybe) 
 
data Trie a = Trie (M.Map a (Trie a)) deriving Show 
 
empty :: Ord a => Trie a 
empty = Trie M.empty 
 
null :: Ord a => Trie a -> Bool 
null (Trie m) = M.null m 

insert :: Ord a => [a] -> Trie a -> Trie a 
insert [] t = t 
insert (x:xs) (Trie m) = Trie m' 
  where m' = M.alter t' x m 
        t' = Just . insert xs . fromMaybe empty 
 
remove :: Ord a => [a] -> Trie a -> Trie a 
remove [] t = t 
remove (x:xs) (Trie m) = Trie m' 
  where m' = M.alter t x m 
        t v = case v of 
              Nothing -> Nothing 
              Just t' -> if null t'' then Nothing 
                                     else Just t'' 
                         where t'' = remove xs t' 

-- contains certain prefix? 
contains :: Ord a => [a] -> Trie a -> Bool 
contains [] _ = True 
contains (x:xs) (Trie m) = case M.lookup x m of 
                           Just t  -> contains xs t 
                           Nothing -> False 
 
-- return subtrie, raise error if not exists that prefix
lookup :: Ord a => [a] -> Trie a -> Trie a 
lookup [] t = t 
lookup (x:xs) t@(Trie m) = lookup xs $ M.findWithDefault t x m