Miscellaneous list operations

ListOperations.hs

import Data.List

-- | What to do when taking the middle of a list where the number of elements
-- around the middle window is odd. Either return the requested number of
-- elements or grow the window by one (making the number of elements on either
-- side of the window the same). The latter is convenient for computing the
-- median of a list (@median = mean . takeMiddle GrowWindow 1@).
data WindowBehavior = KeepWindow  -- ^ @splitMiddle KeepWindow 1 "abcd" = ("a", "b", "cd")@
                    | GrowWindow  -- ^ @splitMiddle GrowWindow 1 "abcd" = ("a", "bc", "d")@
  deriving (Eq, Ord, Bounded, Show, Read)

-- | Drop the @n@ tail elements
--
-- >>> dropTail 3 "hello world"
-- "hello wo"
dropTail :: Int -> [a] -> [a]
dropTail n xs = zipWith const xs (drop n xs)

-- | Take the @n@ tail elements
--
-- >>> takeTail 3 "hello world"
-- "rld"
takeTail :: Int -> [a] -> [a]
takeTail n xs = go xs (drop n xs) where
  go (_:as) (_:bs) = go as bs
  go as     _      = as

-- | Split at the @n@th last element
--
-- >>> splitAtTail 3 "hello world"
-- ("hello wo", "rld")
splitAtTail :: Int -> [a] -> ([a], [a])
splitAtTail n xs = go xs (drop n xs) where
  go (a:as) (_:bs) = let (ls, rs) = go as bs in (a:ls, rs)
  go as     _      = ([], as)

-- | Take the @n@ (or @n+1@, see 'WindowBehavior') middle elements
--
-- >>> takeMiddle KeepWindow 2 "abcd"
-- "bc"
--
-- >>> takeMiddle GrowWindow 2 "abcd"
-- "bc"
--
-- >>> takeMiddle KeepWindow 1 "abcd"
-- "b"
--
-- >>> takeMiddle GrowWindow 1 "abcd"
-- "bc"
takeMiddle :: WindowBehavior -> Int -> [a] -> [a]
-- The first element of the window of size @n@ given a list of length @l@ is at
-- index @(l - n) `div` 2@.
--
-- 'drop' is equivalent to Peano subtraction and iteration in steps of
-- @(_:_:bs)@ is equivalent to Peano division by 2. The second list is used
-- solely as a Peano number.
--
-- Matching the Peano number against @[_]@ vs. @[]@ tells whether the remainder
-- (after the division) is 1 or 0.
takeMiddle behavior n xs = go xs (drop n xs) where
  go (_:as) (_:_:bs) = go as bs
  go as     [_]      | GrowWindow <- behavior
                     = take (n+1) as
  go as     _        = take n     as

-- | Take the @n@ (or @n+1@, see 'WindowBehavior') middle elements as well as
-- the elements preceding and following them
--
-- >>> splitMiddle KeepWindow 2 "abcd"
-- ("a", "bc", "d")
--
-- >>> splitMiddle GrowWindow 2 "abcd"
-- ("a", "bc", "d")
--
-- >>> splitMiddle KeepWindow 1 "abcd"
-- ("a", "b", "cd")
--
-- >>> splitMiddle GrowWindow 1 "abcd"
-- ("a", "bc", "d")
splitMiddle :: WindowBehavior -> Int -> [a] -> ([a], [a], [a])
-- See the comment in 'takeMiddle' regarding the usage of the second list.
splitMiddle behavior n xs = go xs (drop n xs) where
  go (a:as) (_:_:bs) = let (ls, ms, rs) = go as bs in (a:ls, ms, rs)
  go as     [_]      | GrowWindow <- behavior
                     = let (ms, rs) = splitAt (n+1) as in ([], ms, rs)
  go as     _        = let (ms, rs) = splitAt n     as in ([], ms, rs)

-- | Return all windows of @n@ elements to the input list
--
-- >>> windowsOf 3 "hello world"
-- ["hel","ell","llo","lo ","o w"," wo","wor","orl","rld"]
windowsOf :: Int -> [a] -> [[a]]
windowsOf n = dropTail n . map (take n) . tails