func-data-presn/SkewRandList.hs

62 lines
1.5 KiB
Haskell

module SkewRandList
( RList
, cons
, head
, tail
, lookup
, update ) where
import Prelude hiding (head, lookup, tail)
data BTree a = Leaf a
| Node a (BTree a) (BTree a)
deriving (Show, Eq)
data Digit a = Digit !Int (BTree a)
deriving (Show, Eq)
type RList a = [Digit a]
cons :: a -> RList a -> RList a
cons x (Digit w1 t1 : Digit w2 t2 : ts')
| w1 == w2 = Digit (1 + w1 + w2) (Node x t1 t2) : ts'
cons x ts = Digit 1 (Leaf x) : ts
head :: RList a -> a
head (Digit 1 (Leaf x) : _) = x
head (Digit _ (Node x _ _) : _) = x
tail :: RList a -> RList a
tail (Digit 1 (Leaf _) : ts) = ts
tail (Digit w (Node _ t1 t2) : ts) =
Digit w' t1 : Digit w' t2 : ts
where w' = w `div` 2
lookup :: RList a -> Int -> a
lookup (Digit w t : ts) i
| i < w = lookupTree w t i
| otherwise = lookup ts (i - w)
update :: RList a -> Int -> a -> RList a
update (Digit w t : ts) i y
| i < w = Digit w (updateTree w t i y) : ts
| otherwise = Digit w t : update ts (i-w) y
lookupTree :: Int -> BTree a -> Int -> a
lookupTree 1 (Leaf x) 0 = x
lookupTree _ (Node x _ _) 0 = x
lookupTree w (Node _ t1 t2) i
| i < w' = lookupTree w' t1 (i-1)
| otherwise = lookupTree w' t2 (i-1-w')
where w' = w `div` 2
updateTree :: Int -> BTree a -> Int -> a -> BTree a
updateTree 1 (Leaf _) 0 y = Leaf y
updateTree _ (Node _ t1 t2) 0 y = Node y t1 t2
updateTree w (Node x t1 t2) i y
| i < (w`div`2) =
Node x (updateTree w' t1 (i-1) y) t2
| otherwise =
Node x t1 (updateTree w' t2 (i-1-w') y)
where w' = w `div` 2