Programming Praxis: Min Stack

Today’s Programming Praxis task is simple and fun:

Design a data structure that provides push and pop operations, like a stack, plus a third operation that finds the minimum element. All three operations must perform in constant time. You may assume that all elements are distinct.

The solution is to keep an extra stack that keeps track of what the minimum element corresponding to each push was. If the element pushed is smaller than the top element of this internal “min” stack, the new value is pushed, if not the old value is duplicated:

data MinStack a = MinStack {
    stack    :: [a]
  , minStack :: [a]
} deriving (Show)

push :: (Ord a) => a -> MinStack a -> MinStack a
push e (MinStack [] _)      = MinStack [e] [e]
push e (MinStack s (m:ms))  = let s' = e:s
                                  m' = (if m > e then e else m):m:ms
                               in MinStack s' m'

pop :: (Ord a) => MinStack a -> (Maybe a,MinStack a)
pop (MinStack [] _)           = (Nothing,MinStack [] [])
pop (MinStack (e:es) (_:ms))  = (Just e, MinStack es ms)

minElem :: (Ord a) => MinStack a -> Maybe a
minElem (MinStack [] _)    = Nothing
minElem (MinStack _ (m:_)) = Just m

However, this means that we will use twice as much space as we need to store just the elements. Since we are allowed to assume that elements are distinct, the duplication of “min” values is not necessary, instead we can just pop the internal “min” stack when we pop an element that is equal to its top element:

push' :: (Ord a) => a -> MinStack a -> MinStack a
push' e (MinStack [] _)      = MinStack [e] [e]
push' e (MinStack s (m:ms))  = let s' = e:s
                                   m' = (if m > e then e:m:ms else m:ms)
                                in MinStack s' m'

pop' :: (Ord a) => MinStack a -> (Maybe a,MinStack a)
pop' (MinStack [] _)           = (Nothing,MinStack [] [])
pop' (MinStack (e:es) (m:ms))  = (Just e, MinStack es (if e == m then ms else m:ms))

This way we will only increase the size of the internal “min” stack when a new, smallest element is pushed to the stack.