Solution of the ITA's "Strawberry Fields" puzzle.

gistfile1.hs

import Data.Array
import Data.Char
import Data.List
import Data.Maybe
import Data.Ord

import Control.Applicative

import System.Environment
import System.IO.Unsafe

import Data.Hashable
import qualified Data.HashTable.IO as H

-- Ugly (but more performant) memoization using 'unsafePerformIO' as described in:
-- http://augustss.blogspot.com/2011/04/ugly-memoization-heres-problem-that-i.html
-- To further improve performance we use mutable hash tables from the 'hashtables'
-- package instead of Data.Map
memoIO :: (Eq a, Hashable a) => (a -> b) -> IO (a -> IO b)
memoIO f = do
  t <- H.new :: IO (H.CuckooHashTable a b)
  let f' x = do v <- H.lookup t x
                case v of
                  Nothing -> do let { r = f x }; H.insert t x r; return r
                  Just r  -> return r
  return f'

memo :: (Eq a, Hashable a) => (a -> b) -> (a -> b)
memo f = let f' = unsafePerformIO (memoIO f) in \x -> unsafePerformIO (f' x)

data Strawberry = Strawberry { strawberryID :: !Int
                             , strawberryX  :: !Int
                             , strawberryY  :: !Int }

instance Eq Strawberry where
    Strawberry id1 _ _ == Strawberry id2 _ _ = id1 == id2

instance Hashable Strawberry where
    hash = strawberryID

data Greenhouse = Greenhouse !Int !Int !Int !Int deriving Eq

type Field = [Strawberry]
data Cover = Cover [Greenhouse] Cost

type Area = Int
type Cost = Int

strawberries :: Array (Int, Int) Strawberry
strawberries = array ((0, 0), (50, 50)) [((i, j), Strawberry (51 * i + j) i j)
                                             | i <- [0..50]
                                             , j <- [0..50]]

mkStrawberry :: Int -> Int -> Strawberry
mkStrawberry i j = strawberries ! (i, j)

merge :: Cover -> Cover -> Cover
merge (Cover gs1 p1) (Cover gs2 p2) = Cover (gs1 ++ gs2) (p1 + p2)
{-# INLINE merge #-}

area :: Greenhouse -> Area
area (Greenhouse xmin ymin xmax ymax) = (xmax - xmin + 1) * (ymax - ymin + 1)
{-# INLINE area #-}

cost :: Cover -> Cost
cost (Cover _ p) = p
{-# INLINE cost #-}

-- Return the minimum bounding greenhouse of a field.
boundingGreenhouse :: Field -> Greenhouse
boundingGreenhouse ((Strawberry _ x1 y1):ss)
    = foldl extend (Greenhouse x1 y1 x1 y1) ss
    where
      extend (Greenhouse xmin ymin xmax ymax) (Strawberry _ x y)
          = Greenhouse (min xmin x) (min ymin y) (max xmax x) (max ymax y)

-- Cover a field with exactly n greenhouses.  Return Just a cover or
-- Nothing if no such cover exists.  The cover is not guaranteed to
-- be optimal among all possible covers, but it is optimal among the
-- covers that can be obtained by successively subdividing the field
-- into subfields by vertical and horizontal cuts and shrinking the
-- obtained rectangles.
cover' :: (Int, Field) -> Maybe Cover
cover' (n, field) | n > length field = Nothing
cover' (1, field) = Just (Cover [g] p)
    where
      g = boundingGreenhouse field
      p = area g + 10
cover' (n, field) = minimumByCost maybe_covers
    where
      maybe_cover1 = coverSplit strawberryX n field
      maybe_cover2 = coverSplit strawberryY n field
      maybe_covers = [maybe_cover1, maybe_cover2]

minimumByCost :: [Maybe Cover] -> Maybe Cover
minimumByCost maybe_covers
    | null covers
    = Nothing
    | otherwise
    = Just $ minimumBy (comparing cost) covers
    where
      covers = catMaybes maybe_covers

-- Split a field at a given point along the axis specified by the
-- coordinate function coord.  Returns a pair of subfields.
splitField :: (Strawberry -> Int) -> Int -> Field -> (Field, Field)
splitField coord point = partition ((<= point) . coord)

{-

You would think that

init . sort . nub . map coord $ field

cannot be more efficient than the following function that computes
possible split points in linear time (the idea is borrowed from
Bird's "Pearls of Functional Algorithm Design"):

splitPoints :: (Strawberry -> Int) -> Field -> [Int]
splitPoints coord field = tail (indices checklist)
    where
      checklist = accumArray (||) False (0, 50)
                  [(coord strawberry, True) | strawberry <- field]

But at least at the sample inputs the program with the above function
runs 3 (!) times slower (2m6s while the program with the computes the
split points naively runs in under 43s).  Mysteriously, it also uses
more memory.  Apparently, GHC does a pretty good job at fusing the
loops in init . sort . nub . map coord $ field.  Also, our fields are
relatively small, so it may happen that even though splitPoints
requires linear time, the constant factors screw us.

-}

-- Find an optimal cover of a field by n greenhouses by splitting the
-- field along the axis given by the coordinate function coord at all
-- possible places, covering the subfields, and combining the covers.
coverSplit :: (Strawberry -> Int) -> Int -> Field -> Maybe Cover
coverSplit coord n field = minimumByCost maybe_covers
    where
      split_points = init . sort . nub . map coord $ field
      maybe_covers = [liftA2 merge (memoCover' (i,   field1))
                                   (memoCover' (n-i, field2))
                         | i <- [1..n-1], point <- split_points
                         , let (field1, field2) = splitField coord point field]

memoCover' = memo cover'

-- Try to cover a field with at most n greenhouses and choose a cover
-- with the lowest cost.
cover :: Int -> Field -> Cover
cover n field = fromJust $ minimumByCost [memoCover' (i, field) | i <- [1..n]]

-- Parse a field from its matrix representation.
parseField :: [String] -> Field
parseField ls = concat $ zipWith parseLine ls [1..]

-- Parse one line of a matrix representation of the field.
parseLine :: String -> Int -> Field
parseLine line x = [mkStrawberry x y | (c, y) <- zip line [1..], c == '@']

-- Format a cover.  Here (n, m) is the size of the original field.
-- Puts the cost of the cover in the first line.
showCover :: Cover -> (Int, Int) -> String
showCover (Cover gs p) (n, m) = unlines $ show p : [showLine x | x <- [1..n]]
    where
      labelledGreenhouses = zip gs ['A'..]
      showLine x = [showPoint (x, y) | y <- [1..m]]
      isInside (x, y) (Greenhouse xmin ymin xmax ymax, _)
          = xmin <= x && x <= xmax && ymin <= y && y <= ymax
      showPoint (x, y)
          | Just (_, label) <- find (isInside (x,y)) labelledGreenhouses
          = label
          | otherwise
          = '.'

-- A generalization of 'words' and 'lines' allowing to split a list
-- into sublists delimited by the elements satisfying a predicate.
splitBy :: (a -> Bool) -> [a] -> [[a]]
splitBy p xs = case dropWhile p xs of
                 []  -> []
                 xs' -> w : splitBy p xs''
                     where (w, xs'') = break p xs'

isEmptyLine :: String -> Bool
isEmptyLine = all isSpace

-- Process an example.
processExample :: [String] -> IO Cost
processExample (s:ss) = putStrLn (showCover c (rows, cols)) >> return p
    where
      n     = read s :: Int
      rows  = length ss
      cols  = length (head ss)
      field = parseField ss
      c     = cover n field
      p     = cost c

-- Read examples from a file whose name is supplied as a command-line
-- argument, and process each example.
main = do args <- getArgs
          case args of
            [file] -> do text <- readFile file
                         let examples = splitBy isEmptyLine (lines text)
                         costs <- mapM processExample examples
                         let totalCost = sum costs
                         putStrLn $ "Total cost: " ++ show totalCost
            _      -> putStrLn usageMessage
    where
      usageMessage = "usage: strawberry-fields <file>"