Skip to content

Instantly share code, notes, and snippets.

@AndrasKovacs
Last active December 15, 2015 07:19
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save AndrasKovacs/5222282 to your computer and use it in GitHub Desktop.
Save AndrasKovacs/5222282 to your computer and use it in GitHub Desktop.
{-# LANGUAGE TupleSections, RecordWildCards #-}
import Prelude hiding (lookup)
import Control.Arrow (first)
data TTree k v = TTree {val :: !(Maybe v), node :: !(TNode k v)} deriving Show
data TNode k v = Empty | Node {key :: !k, lch, eqch, rch :: !(TTree k v)} deriving Show
empty :: TTree k v
empty = TTree Nothing Empty
modify :: Ord k => ([k] -> TTree k v -> TTree k v) ->
[k] -> TNode k v -> TNode k v
modify f (k:ks) n@(Node {..}) = case compare k key of
LT -> n {lch = f (k:ks) lch}
GT -> n {rch = f (k:ks) rch}
EQ -> n {eqch = f ks eqch}
modify f _ _ = undefined
delete :: Ord k => [k] -> TTree k v -> TTree k v
delete [] t = t {val = Nothing}
delete ks t = t {node = case node t of
Empty -> Empty
n -> modify delete ks n}
update :: Ord k => [k] -> v -> (v -> v) ->
TTree k v -> TTree k v
update [] v f t = t {val = Just $ maybe v f (val t)}
update ks v f t = case node t of
Empty -> insert ks v empty {val = val t}
n -> t { node = modify (\k -> update k v f) ks n }
insert :: Ord k => [k] -> v -> TTree k v -> TTree k v
insert [] v t = t { val = Just v }
insert (k:ks) v t = t {node = case node t of
Empty -> Node k empty (insert ks v empty) empty
n -> modify (flip insert v) (k:ks) n}
lookup :: Ord k => [k] -> TTree k v -> Maybe v
lookup [] t = val t
lookup (k:ks) t = case node t of
Empty -> Nothing
Node {..} -> case compare k key of
LT -> lookup (k:ks) lch
GT -> lookup (k:ks) rch
EQ -> lookup ks eqch
enlist :: Ord k => TTree k v -> [([k], v)]
enlist (TTree {..}) = maybe [] ((:[]).([],)) val ++ case node of
Empty -> []
Node {..} -> concat [enlist lch, map (first (key:)) (enlist eqch), enlist rch]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment