[Haskell-cafe] Re: Red-Blue Stack

[Timothy Goddard]

It won't be O(1) but this is how I would do it. It uses alternating lists of 
red and blue elements. It has to access at most three elements from this list 
for any one operation so as long as we don't have huge blocks of red or blue 
elements performance should be quite good.

The worst case I can think of is if we have an extremely large number of one 
colour followed by a single element of the other then pop that single element 
off the stack. This would require two lists (before and after the single 
element) to be combined with ++, taking time linear to the size of the first 
list.

Anyway, here's some code:

module RBStack
  where

data RBColour = Red | Blue
  deriving (Show, Eq)

data RBStack a =  RBStack {
  headColour :: RBColour,
  stackElems :: [[a]]
}
  deriving (Show, Eq)

otherCol :: RBColour -> RBColour
otherCol Red = Blue
otherCol Blue = Red

empty :: RBStack a
empty = RBStack Red []

push :: RBColour -> a -> RBStack a -> RBStack a
push col val stack
  | null (stackElems stack) = RBStack col [[val]]
  | headColour stack == col = RBStack col ((val:e):es)
  | otherwise = RBStack col ([val]:e:es)
  where
    (e:es) = stackElems stack

popColour :: RBColour -> RBStack a -> (Maybe a, RBStack a)
popColour col stack
  | null (stackElems stack) = (Nothing, stack)
  | headColour stack == col = (Just (head e), if null (tail e)
      then (RBStack (otherCol col) es)
      else (RBStack col ((tail e):es)))
  | otherwise = if null es
      then (Nothing, empty)
      else let (f:fs) = es in (Just (head f), if null (tail f)
        then (if null fs then (RBStack (otherCol col) [e]) else (RBStack 
(otherCol col) ((e ++ (head fs)):(tail fs))))
        else RBStack (otherCol col) (e:(tail f):fs))
  where
    (e:es) = stackElems stack

pop :: RBStack a -> (Maybe (RBColour, a), RBStack a)
pop stack
  | null (stackElems stack) = (Nothing, stack)
  | otherwise = (Just (col, head e), if null (tail e) then (RBStack (otherCol 
col) es) else (RBStack col ((tail e):es)))
  where
    (e:es) = stackElems stack
    col = headColour stack

peek :: RBStack a -> Maybe (RBColour, a)
peek = fst . pop