Posts Tagged ‘flag’

Programming Praxis – Dutch National Flag

March 5, 2013

In today’s Programming Praxis exercise, our goal is to implement a sorting algorithm for lists with three different elements that works in linear time. Let’s get started, shall we?

import qualified Data.Vector as V

We’re only allowed to use index and swap operations. Swap isn’t defined in the Vector package, but is easy to express as a bulk update.

swap :: Int -> Int -> V.Vector a -> V.Vector a
swap i j a = a V.// [(i,a V.! j), (j,a V.! i)]

To sort the array, we keep track of where the next red and blue elements should be inserted, as well as the current index. When encountering red and blue elements, they are shifted to the correct location. If the element was blue, we have to test again since it could be anything. If it was red, the element that’s swapped to the current location will always be white, so we can move on. Once we reach the start of the group of blue elements at the end we can stop.

flag :: V.Vector Char -> V.Vector Char
flag xs = f (0, V.length xs - 1) xs 0 where
  f (r,b) a n = if n > b then a else case a V.! n of
    'R' -> f (r+1,b  ) (swap n r a) (n+1)
    'B' -> f (r,  b-1) (swap n b a) n
    _   -> f (r,  b  ) a            (n+1)

Some tests to see if everything is working properly:

test :: String -> Bool
test x = flag (V.fromList x) == V.fromList
  (filter (== 'R') x ++ filter (== 'W') x ++ filter (== 'B') x)

main :: IO ()
main = do print $ test ""
          print $ test "W"
          print $ test "R"
          print $ test "B"
          print $ test "RWB"
          print $ test "BWR"
          print $ test "RWBR"