49 lines
1.5 KiB
Haskell
49 lines
1.5 KiB
Haskell
{-# LANGUAGE TemplateHaskell #-}
|
|
|
|
module SimpleTrie where
|
|
|
|
import Control.Lens
|
|
import qualified Data.Map.Strict as M
|
|
|
|
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
|