43 lines
1.5 KiB
Haskell
43 lines
1.5 KiB
Haskell
{-# LANGUAGE TemplateHaskell #-}
|
|
module SimpleTrie where
|
|
import qualified Data.Map.Strict as M
|
|
import Control.Lens
|
|
data Trie a b = Trie
|
|
{ _trieMVal :: Maybe b
|
|
, _trieChildren :: M.Map a (Trie a b)
|
|
}
|
|
makeLenses ''Trie
|
|
|
|
emptyTrie :: Trie a b
|
|
emptyTrie = Trie Nothing M.empty
|
|
|
|
singletonTrie :: Ord a => [a] -> b -> Trie a b
|
|
singletonTrie (k:ks) x = Trie Nothing $ M.singleton k $ singletonTrie ks x
|
|
singletonTrie [] x = Trie (Just x) M.empty
|
|
|
|
insertInTrie :: Ord a => [a] -> b -> Trie a b -> Trie a b
|
|
insertInTrie [] x = trieMVal ?~ x
|
|
insertInTrie (k:ks) x = trieChildren %~ M.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 M.!? k >>= lookupTrie ks
|
|
lookupTrie [] t = _trieMVal t
|
|
|
|
multiLookupTrie :: Ord a => [a] -> Trie a b -> [([a],b)]
|
|
multiLookupTrie xs t@(Trie my ch)
|
|
| null xs || null ch = maybe [] (\y -> [([],y)]) my
|
|
| otherwise = let (z:zs) = xs in
|
|
case M.lookup z ch of
|
|
Nothing -> multiLookupTrie zs t
|
|
Just t' -> map (_1 %~ (z:)) (multiLookupTrie zs t') ++ multiLookupTrie zs t
|
|
|
|
multiLookupTrieI :: Ord a => [(a,c)] -> Trie a b -> [([c],b)]
|
|
multiLookupTrieI xs t@(Trie my ch)
|
|
| null xs || null ch = maybe [] (\y -> [([],y)]) my
|
|
| otherwise = let ((z1,z2):zs) = xs in
|
|
case M.lookup z1 ch of
|
|
Nothing -> multiLookupTrieI zs t
|
|
Just t' -> map (_1 %~ (z2:)) (multiLookupTrieI zs t') ++ multiLookupTrieI zs t
|