Programming Praxis – An Array Of Two Symbols

In today’s exercise, our goal is to write a function that, given a list consisting of m of one symbol followed by n of another symbol, returns the starting index of the second group of symbols in O(log m) time. Let’s get started, shall we?

import qualified Data.Vector as V
import Test.QuickCheck

The basis idea is simple: starting at the first element, keep doubling the index until we encounter a symbol from the second group or we run past the end of the array. Repeat the process, only now looking at the subarray from the last m to the first n (or the end of the array). When this array has only two values (by definition an m and an n), we’ve found our answer.

Unlike the provided solution, we don’t use a binary search once we’ve found our bounds. We already have one O(log m) test, so why bother writing another?

search :: V.Vector Char -> Int
search xs = f 0 (V.length xs - 1) where
    f start end = case span ((xs V.! 0 ==) . (xs V.!)) . takeWhile (<= end) .
                       map (start +) $ 0 : iterate (*2) 1
                  of   ([_],[n]) -> n
                       (ms ,[] ) -> f (last ms) end
                       (ms ,n:_) -> f (last ms) n

To test if everything is working correctly, we use a few manual tests and quickCheck for an automated one.

test :: Property
test = forAll (choose (1,100)) $ \i ->
       forAll (choose (1,100)) $ \j ->
       search (V.fromList $ replicate i 'm' ++ replicate j 'n') == i

main :: IO ()
main = do print $ search (V.fromList "mn") == 1
          print $ search (V.fromList "mmmnn") == 3
          print $ search (V.fromList "mmmmmmnnnnnnnnnnn") == 6
          quickCheck test