AVL木によるMap

AVL.hs

{-# LANGUAGE DeriveFunctor #-}
module Data.AVL
  ( AVL
  , size
  , toList
  , fromList
  , empty
  , singleton
  , findMin
  , findMax
  , lookup
  , insert
  , delete
  , deleteMin
  , deleteMax
  )
  where

import Prelude hiding (lookup)
import Control.Applicative hiding (empty)

data AVL k a = Leaf | Node !Int k a (AVL k a) (AVL k a)
  deriving (Eq,Functor)

instance (Show k, Show a) => Show (AVL k a) where
  showsPrec p t = showParen (p>10) $ ("fromList " ++) . shows (toList t)

fromList :: Ord k => [(k,a)] -> AVL k a
fromList = foldr (uncurry insert) empty

toList :: AVL k a -> [(k,a)]
toList t = go t []
  where
    go Leaf = id
    go (Node _ k a l r) = go l . ((k,a) :) . go r

showTree :: (Show k, Show a) => AVL k a -> String
showTree = unlines . go
  where
    showNode = ("+- " ++)
    indentl (s:ss) = showNode s : map ("|  " ++) ss
    indentr (s:ss) = showNode s : map ("   " ++) ss
    go Leaf = [""]
    go (Node _ k a l r) = show (k,a) : indentl (go l) ++ indentr (go r)

findMin :: Ord k => AVL k a -> Maybe (k,a)
findMin Leaf = Nothing
findMin (Node _ k a Leaf _) = Just (k,a)
findMin (Node _ _ _ l _) = findMin l

findMax :: Ord k => AVL k a -> Maybe (k,a)
findMax Leaf = Nothing
findMax (Node _ k a _ Leaf) = Just (k,a)
findMax (Node _ _ _ _ r) = findMax r

size :: AVL k a -> Int
size Leaf = 0
size (Node _ _ _ l r) = 1 + size l + size r

empty :: AVL k a
empty = Leaf

singleton :: k -> a -> AVL k a
singleton k a = node k a empty empty

height :: AVL k a -> Int
height Leaf = 0
height (Node i _ _ _ _) = i

nodeHeight :: AVL k a -> AVL k a -> Int
nodeHeight l r = 1 + max (height l) (height r)

lookup :: Ord k => k -> AVL k a -> Maybe a
lookup ka = go
  where
    go Leaf = Nothing
    go (Node _ kb a l r)
      | ka < kb = go l
      | ka > kb = go r
      | otherwise = Just a

node :: k -> a -> AVL k a -> AVL k a -> AVL k a
node k a l r = Node (nodeHeight l r) k a l r

rotateL1, rotateL2, rotateR1, rotateR2 :: AVL k a -> AVL k a
rotateL1 (Node _ ka a l (Node _ kb b rl rr)) = node kb b (node ka a l rl) rr
rotateL2 (Node _ ka a l (Node _ kb b (Node _ kc c rll rlr) rr))
  = node kc c (node ka a l rll) (node kb b rlr rr)
rotateR1 (Node _ ka a (Node _ kb b ll lr) r) = node kb b ll (node ka a lr r)
rotateR2 (Node _ ka a (Node _ kb b ll (Node _ kc c lrl lrr)) r)
  = node kc c (node kb b ll lrl) (node ka a lrr r)

label :: AVL k a -> Ordering
label Leaf = EQ
label (Node _ _ _ l r) = compare (height l) (height r)

insert :: Ord k => k -> a -> AVL k a -> AVL k a
insert ka a = go
  where
    go Leaf = singleton ka a
    go t@(Node _ kb b l r)
      | ka < kb = let l' = go l
                      t' = node kb b l' r in
        case (lab, label l') of
          (GT, GT) -> rotateR1 t'
          (GT, LT) -> rotateR2 t'
          _  -> t'
      | ka > kb = let r' = go r
                      t' = node kb b l r' in
        case (lab, label r') of
          (LT, LT) -> rotateL1 t'
          (LT, GT) -> rotateL2 t'
          _  -> t'
      | otherwise = t
      where
        lab = label t

deleteMin :: Ord k => AVL k a -> Maybe (k, a, AVL k a)
deleteMin Leaf = Nothing
deleteMin (Node _ k a Leaf _) = Just (k, a, Leaf)
deleteMin (Node _ k a l r) = Just (km, m, node k a l' r)
  where
    Just (km, m, l') = deleteMin l

deleteMax :: Ord k => AVL k a -> Maybe (k, a, AVL k a)
deleteMax Leaf = Nothing
deleteMax (Node _ k a _ Leaf) = Just (k, a, Leaf)
deleteMax (Node _ k a l r) = Just (km, m, node k a l r')
  where
    Just (km, m, r') = deleteMax r

delete :: Ord k => k -> AVL k a -> AVL k a
delete ka = go
  where
    go Leaf = Leaf
    go t@(Node _ kb b l r)
      | ka < kb = let l' = go l
                      t' = node kb b l' r in
        case (lab, label l') of
          (LT, LT) -> rotateL2 t'
          (LT, _ ) -> rotateL1 t'
          _  -> t'
      | ka > kb = let r' = go r
                      t' = node kb b l r' in
        case (lab, label r') of
          (GT, GT) -> rotateR2 t'
          (GT, _ ) -> rotateR1 t'
          _  -> t'
      | size l < size r = case deleteMin r of
        Nothing -> l
        Just (km,m,r') -> node km m l r'
      | otherwise = case deleteMax l of
        Nothing -> r
        Just (km,m,l') -> node km m l' r
      where
        lab = label t

key :: AVL k a -> Maybe k
key Leaf = Nothing
key (Node _ k _ _ _) = Just k

compareAVL :: AVL k a -> AVL k a -> Ordering
compareAVL l r
  | hl + 1 < hr = LT
  | hl > hr + 1 = GT
  | otherwise = EQ
  where
    hl = height l
    hr = height r

check :: Ord k => AVL k a -> Bool
check Leaf = True
check (Node i k _ l r)
  =  i == nodeHeight l r
  && compareAVL l r == EQ
  && maybeBool ((< k) <$> key l)
  && maybeBool ((> k) <$> key r)
  && check l
  && check r
  where
    maybeBool = maybe True id