Timothy Goddard tim at goddard.net.nz
Mon Sep 29 19:12:43 EDT 2008

```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 {
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