27 lines
788 B
Haskell
27 lines
788 B
Haskell
{-# LANGUAGE TemplateHaskell #-}
|
|
module SimpleTrie where
|
|
import Data.Map.Strict
|
|
import Control.Lens
|
|
data Trie a b = Trie
|
|
{ _trieMVal :: Maybe b
|
|
, _trieChildren :: Map a (Trie a b)
|
|
}
|
|
makeLenses ''Trie
|
|
|
|
emptyTrie :: Trie a b
|
|
emptyTrie = Trie Nothing empty
|
|
|
|
singletonTrie :: Ord a => [a] -> b -> Trie a b
|
|
singletonTrie (k:ks) x = Trie Nothing $ singleton k $ singletonTrie ks x
|
|
singletonTrie [] x = Trie (Just x) empty
|
|
|
|
insertInTrie :: Ord a => [a] -> b -> Trie a b -> Trie a b
|
|
insertInTrie [] x = trieMVal ?~ x
|
|
insertInTrie (k:ks) x = trieChildren %~ insertWith f k (singletonTrie ks x)
|
|
where
|
|
f _ = insertInTrie ks x
|
|
|
|
lookupTrie :: Ord a => [a] -> Trie a b -> Maybe b
|
|
lookupTrie (k:ks) t = _trieChildren t !? k >>= lookupTrie ks
|
|
lookupTrie [] t = _trieMVal t
|