74 lines
2.5 KiB
Haskell
74 lines
2.5 KiB
Haskell
{-# OPTIONS -Wno-incomplete-uni-patterns #-}
|
|
{-# LANGUAGE TemplateHaskell #-}
|
|
|
|
module SimpleTrie where
|
|
|
|
import Data.Monoid
|
|
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
|
|
|
|
----nextTrie :: Ord a => [a] -> Trie a b -> Maybe (b, Trie a b)
|
|
--nextTrie :: [a] -> Trie a b -> Maybe (b, Trie a b)
|
|
--nextTrie (k:ks) t = undefined
|
|
--nextTrie [] t = Nothing
|
|
|
|
splitLookupTrie :: Ord a => [a] -> Trie a b -> (Trie a b,Maybe b,Trie a b)
|
|
splitLookupTrie (k:ks) (Trie x xs) = case m of
|
|
Nothing -> (Trie x l,Nothing,Trie x r)
|
|
Just t' -> let (l',m',r') = splitLookupTrie ks t'
|
|
in (Trie x (M.insert k l' l),m', Trie x (M.insert k r' r))
|
|
where
|
|
(l,m,r) = M.splitLookup k xs
|
|
splitLookupTrie [] (Trie x xs) = (Trie Nothing mempty,x, Trie Nothing xs)
|
|
-- the above would be better if it removed empty trie strings; the left trie
|
|
-- will always be empty and the right trie might be empty.
|
|
-- I cannot think of a smart way to do this
|
|
|
|
firstTrie :: Ord a => Trie a b -> Maybe b
|
|
firstTrie t = case t ^. trieMVal of
|
|
Just x -> Just x
|
|
Nothing -> getFirst $ foldMap (First . firstTrie) (t ^. trieChildren)
|
|
|
|
-- # OPTIONS -Wno-incomplete-uni-patterns #-}
|
|
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
|