AvlTree-2.4: Data/Tree/AVL/Size.hs
{-# OPTIONS_GHC -fglasgow-exts #-}
-----------------------------------------------------------------------------
-- |
-- Module : Data.Tree.AVL.Size
-- Copyright : (c) Adrian Hey 2004,2005
-- License : BSD3
--
-- Maintainer : http://homepages.nildram.co.uk/~ahey/em.png
-- Stability : stable
-- Portability : portable
--
-- AVL Tree size related utilities.
-----------------------------------------------------------------------------
module Data.Tree.AVL.Size
(-- * AVL tree size utilities.
size,addSize,fastAddSize,clipSize
) where
import Data.Tree.AVL.Types(AVL(..))
import Data.Tree.AVL.Internals.HeightUtils(addHeight)
#ifdef __GLASGOW_HASKELL__
import GHC.Base
#include "ghcdefs.h"
#else
#include "h98defs.h"
#endif
-- | Counts the total number of elements in an AVL tree.
--
-- @'size' = 'addSize' 0@
--
-- Complexity: O(n)
{-# INLINE size #-}
size :: AVL e -> Int
size = addSize 0
-- | Adds the size of a tree to the first argument.
-- This is just a convenience wrapper for 'fastAddSize'.
--
-- Complexity: O(n)
{-# INLINE addSize #-}
addSize :: Int -> AVL e -> Int
addSize ASINT(n) t = ASINT(fastAddSize n t)
{-----------------------------------------
Notes for fast size calculation.
case (h,avl)
(0,_ ) -> 0 -- Must be E
(1,_ ) -> 1 -- Must be (Z E _ E )
(2,N _ _ _) -> 2 -- Must be (N E _ (Z E _ E))
(2,Z _ _ _) -> 3 -- Must be (Z (Z E _ E) _ (Z E _ E))
(2,P _ _ _) -> 2 -- Must be (P (Z E _ E) _ E )
(3,N _ _ r) -> 2 + size 2 r -- Must be (N (Z E _ E) _ r )
(3,P l _ _) -> 2 + size 2 l -- Must be (P l _ (Z E _ E))
------------------------------------------}
-- | Fast algorithm to calculate size. This avoids visiting about 50% of tree nodes
-- by using fact that trees with small heights can only have particular shapes.
-- So it's still O(n), but with substantial saving in constant factors.
--
-- Complexity: O(n)
fastAddSize :: UINT -> AVL e -> UINT
fastAddSize n E = n
fastAddSize n (N l _ r) = case addHeight L(2) l of
L(2) -> INCINT2(n)
L(3) -> fas2 INCINT2(n) r
h -> fasNP n h l r
fastAddSize n (Z l _ r) = case addHeight L(1) l of
L(1) -> INCINT1(n)
L(2) -> INCINT3(n)
L(3) -> fas2 (fas2 INCINT1(n) l) r
h -> fasZ n h l r
fastAddSize n (P l _ r) = case addHeight L(2) r of
L(2) -> INCINT2(n)
L(3) -> fas2 INCINT2(n) l
h -> fasNP n h r l
-- Parent Height (h) >= 4 !!
fasNP,fasZ :: UINT -> UINT -> AVL e -> AVL e -> UINT
fasNP n h l r = fasG3 (fasG2 INCINT1(n) DECINT2(h) l) DECINT1(h) r
fasZ n h l r = fasG3 (fasG3 INCINT1(n) DECINT1(h) l) DECINT1(h) r
-- h>=2 !!
fasG2 :: UINT -> UINT -> AVL e -> UINT
fasG2 n L(2) t = fas2 n t
fasG2 n h t = fasG3 n h t
{-# INLINE fasG2 #-}
-- h>=3 !!
fasG3 :: UINT -> UINT -> AVL e -> UINT
fasG3 n L(3) (N _ _ r) = fas2 INCINT2(n) r
fasG3 n L(3) (Z l _ r) = fas2 (fas2 INCINT1(n) l) r
fasG3 n L(3) (P l _ _) = fas2 INCINT2(n) l
fasG3 n h (N l _ r) = fasNP n h l r -- h>=4
fasG3 n h (Z l _ r) = fasZ n h l r -- h>=4
fasG3 n h (P l _ r) = fasNP n h r l -- h>=4
fasG3 _ _ E = error "fastAddSize: Bad Tree." -- impossible
-- h=2 !!
fas2 :: UINT -> AVL e -> UINT
fas2 n (N _ _ _) = INCINT2(n)
fas2 n (Z _ _ _) = INCINT3(n)
fas2 n (P _ _ _) = INCINT2(n)
fas2 _ E = error "fastAddSize: Bad Tree." -- impossible
{-# INLINE fas2 #-}
-----------------------------------------------------------------------
----------------------- fastAddSize Ends Here -------------------------
-----------------------------------------------------------------------
-- | Returns the exact tree size in the form @('Just' n)@ if this is less than or
-- equal to the input clip value. Returns @'Nothing'@ of the size is greater than
-- the clip value. This function exploits the same optimisation as 'fastAddSize'.
--
-- Complexity: O(min n c) where n is tree size and c is clip value.
clipSize :: Int -> AVL e -> Maybe Int
clipSize ASINT(c) t = let c_ = cSzh c t in if c_ LTN L(0)
then Nothing
else Just ASINT(SUBINT(c,c_))
-- First entry calculates initial height
cSzh :: UINT -> AVL e -> UINT
cSzh c E = c
cSzh c (N l _ r) = case addHeight L(2) l of
L(2) -> DECINT2(c)
L(3) -> cSzNP3 c r
h -> cSzNP c h l r
cSzh c (Z l _ r) = case addHeight L(1) l of
L(1) -> DECINT1(c)
L(2) -> DECINT3(c)
L(3) -> cSzZ3 c l r
h -> cSzZ c h l r
cSzh c (P l _ r) = case addHeight L(2) r of
L(2) -> DECINT2(c)
L(3) -> cSzNP3 c l
h -> cSzNP c h r l
-- Parent Height = 3 !!
cSzNP3 :: UINT -> AVL e -> UINT
cSzNP3 c t = if c LTN L(4) then L(-1) else cSz2 DECINT2(c) t
cSzZ3 :: UINT -> AVL e -> AVL e -> UINT
cSzZ3 c l r = if c LTN L(5) then L(-1)
else let c_ = cSz2 DECINT1(c) l
in if c_ LTN L(2) then L(-1)
else cSz2 c_ r
-- Parent Height (h) >= 4 !!
cSzNP,cSzZ :: UINT -> UINT -> AVL e -> AVL e -> UINT
cSzNP c h l r = if c LTN L(7) then L(-1)
else let c_ = cSzG2 DECINT1(c) DECINT2(h) l -- (h-2) >= 2
in if c_ LTN L(4) then L(-1)
else cSzG3 c_ DECINT1(h) r -- (h-1) >= 3
cSzZ c h l r = if c LTN L(9) then L(-1)
else let c_ = cSzG3 DECINT1(c) DECINT1(h) l -- (h-1) >= 3
in if c_ LTN L(4) then L(-1)
else cSzG3 c_ DECINT1(h) r -- (h-1) >= 3
-- h>=2 !!
cSzG2 :: UINT -> UINT -> AVL e -> UINT
cSzG2 c L(2) t = cSz2 c t
cSzG2 c h t = cSzG3 c h t
{-# INLINE cSzG2 #-}
-- h>=3 !!
cSzG3 :: UINT -> UINT -> AVL e -> UINT
cSzG3 c L(3) (N _ _ r) = cSzNP3 c r
cSzG3 c L(3) (Z l _ r) = cSzZ3 c l r
cSzG3 c L(3) (P l _ _) = cSzNP3 c l
cSzG3 c h (N l _ r) = cSzNP c h l r -- h>=4
cSzG3 c h (Z l _ r) = cSzZ c h l r -- h>=4
cSzG3 c h (P l _ r) = cSzNP c h r l -- h>=4
cSzG3 _ _ E = error "clipSize: Bad Tree." -- impossible
-- h=2 !!
cSz2 :: UINT -> AVL e -> UINT
cSz2 c (N _ _ _) = DECINT2(c)
cSz2 c (Z _ _ _) = DECINT3(c)
cSz2 c (P _ _ _) = DECINT2(c)
cSz2 _ E = error "clipSize: Bad Tree." -- impossible
{-# INLINE cSz2 #-}
-----------------------------------------------------------------------
------------------------- clipSize Ends Here --------------------------
-----------------------------------------------------------------------