summaryrefslogtreecommitdiff
path: root/Data/Tree/Splay.hs
diff options
context:
space:
mode:
Diffstat (limited to 'Data/Tree/Splay.hs')
-rw-r--r--Data/Tree/Splay.hs139
1 files changed, 100 insertions, 39 deletions
diff --git a/Data/Tree/Splay.hs b/Data/Tree/Splay.hs
index 496ffac..2b096eb 100644
--- a/Data/Tree/Splay.hs
+++ b/Data/Tree/Splay.hs
@@ -1,21 +1,21 @@
--
--- Copyright (c) 2009 Brendan Hickey - http://bhickey.net
--- Simplified BSD License (see http://www.opensource.org/licenses/bsd-license.php)
+-- Copyright (c) 2009 - 2010 Brendan Hickey - http://bhickey.net
+-- New BSD License (see http://www.opensource.org/licenses/bsd-license.php)
--
module Data.Tree.Splay
-(SplayTree, head, tail, singleton, empty, null, fromList, fromAscList, toList, toAscList, insert, lookup)
+(SplayTree, head, tail, singleton, empty, null, fromList, fromAscList, toList, toAscList, insert, lookup, (!!), splay, size, delete)
where
-import Prelude hiding (head, tail, lookup, null)
+import Prelude hiding (head, tail, lookup, null, (!!))
data (Ord k) => SplayTree k v =
Leaf
- | SplayTree k v (SplayTree k v) (SplayTree k v) deriving (Ord, Eq)
+ | SplayTree k v Int (SplayTree k v) (SplayTree k v) deriving (Ord, Eq)
-- | /O(1)/. 'singleton' constructs a splay tree containing one element.
singleton :: (Ord k) => (k,v) -> SplayTree k v
-singleton (k,v) = SplayTree k v Leaf Leaf
+singleton (k,v) = SplayTree k v 0 Leaf Leaf
-- | /O(1)/. 'empty' constructs an empty splay tree.
empty :: (Ord k) => SplayTree k v
@@ -26,61 +26,122 @@ null :: (Ord k) => SplayTree k v -> Bool
null Leaf = True
null _ = False
+size :: (Ord k) => SplayTree k v -> Int
+size Leaf = 0
+size (SplayTree _ _ d _ _) = d
-- | /Amortized O(lg n)/. Given a splay tree and a key, 'lookup' attempts to find a node with the specified key and splays this node to the root. If the key is not found, the nearest node is brought to the root of the tree.
-lookup :: (Ord k) => SplayTree k v -> k -> SplayTree k v
-lookup Leaf _ = Leaf
-lookup n@(SplayTree k v l r) sk =
- if sk == k
- then n
- else if k > sk
- then case lookup l sk of
- Leaf -> n
- (SplayTree k1 v1 l1 r1) -> (SplayTree k1 v1 l1 (SplayTree k v r1 r))
- else case lookup r sk of
- Leaf -> n
- (SplayTree k1 v1 l1 r1) -> (SplayTree k1 v1 (SplayTree k v l l1) r1)
-
--- | /Amortized O(lg n)/. Given a splay tree and a key-value pair, 'insert' places the the pair into the tree in BST order.
-insert :: (Ord k) => SplayTree k v -> (k,v) -> SplayTree k v
-insert t (k,v) =
- case lookup t k of
- Leaf -> (SplayTree k v Leaf Leaf)
- (SplayTree k1 v1 l r) ->
+lookup :: (Ord k) => k -> SplayTree k v -> SplayTree k v
+lookup _ Leaf = Leaf
+lookup k' t@(SplayTree k _ _ l r)
+ | k' == k = t
+ | k > k' =
+ case lookup k' l of
+ Leaf -> t
+ lt -> zig lt t
+ | otherwise =
+ case lookup k' r of
+ Leaf -> t
+ rt -> zag t rt
+
+-- | Locates the i^{th} element in BST order without splaying it.
+(!!) :: (Ord k) => SplayTree k v -> Int -> (k,v)
+(!!) Leaf _ = error "index out of bounds"
+(!!) (SplayTree k v d l r) n =
+ if n > d
+ then error "index out of bounds"
+ else
+ let l' = size l in
+ if n == l'
+ then (k,v)
+ else if n <= l'
+ then l !! n
+ else r !! (n - l')
+
+-- | Splays the i^{th} element in BST order
+splay :: (Ord k) => SplayTree k v -> Int -> SplayTree k v
+splay Leaf _ = error "index out of bounds"
+splay t@(SplayTree _ _ d l r) n =
+ if n > d
+ then error "index out of bounds"
+ else
+ let l' = size l in
+ if n == l'
+ then t
+ else if n <= l'
+ then case splay l n of
+ Leaf -> error "index out of bounds"
+ lt -> zig lt t
+ else case splay r (n - l') of
+ Leaf -> error "index out of bounds"
+ rt -> zag t rt
+
+-- | /O(1)/. zig rotates its first argument up
+zig :: (Ord k) => SplayTree k v -> SplayTree k v -> SplayTree k v
+zig Leaf _ = error "tree corruption"
+zig _ Leaf = error "tree corruption"
+zig (SplayTree k1 v1 _ l1 r1) (SplayTree k v d _ r) =
+ SplayTree k1 v1 d l1 (SplayTree k v (d - size l1 - 1) r1 r)
+
+-- | /O(1)/. zig rotates its second argument up
+zag :: (Ord k) => SplayTree k v -> SplayTree k v -> SplayTree k v
+zag Leaf _ = error "tree corruption"
+zag _ Leaf = error "tree corruption"
+zag (SplayTree k v d l _) (SplayTree k1 v1 _ l1 r1) =
+ SplayTree k1 v1 d (SplayTree k v (d - size r1 - 1) l l1) r1
+
+-- | /Amortized O(lg n)/. Given a splay tree and a key-value pair, 'insert' places the the pair into the tree in BST order. This function is unsatisfying.
+insert :: (Ord k) => k -> v -> SplayTree k v -> SplayTree k v
+insert k v t =
+ case lookup k t of
+ Leaf -> (SplayTree k v 0 Leaf Leaf)
+ (SplayTree k1 v1 d l r) ->
if k1 < k
- then (SplayTree k v (SplayTree k1 v1 l Leaf) r)
- else (SplayTree k v l (SplayTree k1 v1 Leaf r))
-
+ then SplayTree k v (d + 1) (SplayTree k1 v1 (d - size r + 1) l Leaf) r
+ else SplayTree k v (d + 1) l (SplayTree k1 v1 (d - size l + 1) Leaf r)
-- | /O(1)/. 'head' returns the key-value pair of the root.
head :: (Ord k) => SplayTree k v -> (k,v)
head Leaf = error "head of empty tree"
-head (SplayTree k v _ _) = (k,v)
+head (SplayTree k v _ _ _) = (k,v)
-- | /Amortized O(lg n)/. 'tail' removes the root of the tree and merges its subtrees
tail :: (Ord k) => SplayTree k v -> SplayTree k v
tail Leaf = error "tail of empty tree"
-tail (SplayTree _ _ Leaf r) = r
-tail (SplayTree _ _ l Leaf) = l
-tail (SplayTree _ _ l r) =
+tail (SplayTree _ _ _ Leaf r) = r
+tail (SplayTree _ _ _ l Leaf) = l
+tail (SplayTree _ _ _ l r) =
case splayRight l of
- (SplayTree k v l1 Leaf) -> (SplayTree k v l1 r)
+ (SplayTree k v d l1 Leaf) -> (SplayTree k v (d + size r) l1 r)
_ -> error "splay tree corruption"
+delete :: (Ord k) => k -> SplayTree k v -> SplayTree k v
+delete _ Leaf = Leaf
+delete k t =
+ case lookup k t of
+ t'@(SplayTree k1 _ _ _ _) ->
+ if k == k1
+ then tail t'
+ else t'
+ Leaf -> error "splay tree corruption"
+
+
splayRight :: (Ord k) => SplayTree k v -> SplayTree k v
splayRight Leaf = Leaf
-splayRight h@(SplayTree _ _ _ Leaf) = h
-splayRight (SplayTree k1 v1 l1 (SplayTree k2 v2 l2 r2)) = splayRight $ (SplayTree k2 v2 (SplayTree k1 v1 l1 l2) r2)
+splayRight h@(SplayTree _ _ _ _ Leaf) = h
+splayRight (SplayTree k1 v1 d1 l1 (SplayTree k2 v2 _ l2 r2)) =
+ splayRight (SplayTree k2 v2 d1 (SplayTree k1 v1 (d1 - size r2) l1 l2) r2)
splayLeft :: (Ord k) => SplayTree k v -> SplayTree k v
splayLeft Leaf = Leaf
-splayLeft h@(SplayTree _ _ Leaf _) = h
-splayLeft (SplayTree k1 v1 (SplayTree k2 v2 l2 r2) r1) = splayLeft $ (SplayTree k2 v2 l2 (SplayTree k1 v1 r2 r1))
+splayLeft h@(SplayTree _ _ _ Leaf _) = h
+splayLeft (SplayTree k1 v1 d1 (SplayTree k2 v2 _ l2 r2) r1) =
+ splayLeft (SplayTree k2 v2 d1 l2 (SplayTree k1 v1 (d1 - size l2) r2 r1))
-- | /O(n lg n)/. Constructs a splay tree from an unsorted list of key-value pairs.
fromList :: (Ord k) => [(k,v)] -> SplayTree k v
fromList [] = Leaf
-fromList l = foldl (\ acc x -> insert acc x) Leaf l
+fromList l = foldl (\ acc (k,v) -> insert k v acc) Leaf l
-- | /O(n lg n)/. Constructs a splay tree from a list of key-value pairs sorted in ascending order.
fromAscList :: (Ord k) => [(k,v)] -> SplayTree k v
@@ -92,6 +153,6 @@ toList = toAscList
-- | /O(n lg n)/. 'toAscList' converts a splay tree to a list of key-value pairs sorted in ascending order.
toAscList :: (Ord k) => SplayTree k v -> [(k,v)]
-toAscList h@(SplayTree _ _ Leaf _) = (head h):(toAscList $ tail h)
+toAscList h@(SplayTree _ _ _ Leaf _) = head h : toAscList (tail h)
toAscList Leaf = []
toAscList h = toAscList $ splayLeft h