Smash The Code Contest -- Initial commit, Silver league

SmashTheCode/SmashTheCode.hs

@@ -0,0 +1,160 @@
+{-# LANGUAGE FlexibleContexts, LambdaCase, TupleSections #-}
+
+import Control.Applicative
+import Control.Arrow (first, second, (&&&))
+import Control.Monad
+import Control.Monad.State
+import Control.Monad.Writer
+import Data.Function
+import Data.List
+import Data.Maybe
+import Data.Monoid
+import Data.Array (Array, (!), (//))
+import Data.Set (Set)
+import System.IO
+
+import qualified Data.Array       as A
+import qualified Data.Foldable    as F
+import qualified Data.Traversable as T
+import qualified Data.Set         as S
+
+{-# ANN module "HLint: ignore Use if" #-}
+{-# ANN module "HLint: ignore Redundant $" #-}
+{-# ANN module "HLint: ignore Redundant do" #-}
+
+type Color    = Int
+type Rotation = Int
+type Block    = (Color, Color)
+type Column   = Int
+type Row      = Int
+data Cell     = Empty | Skull | Color Color deriving (Eq, Ord, Show)
+type Grid     = Array (Column, Row) Cell
+
+main :: IO ()
+main = do
+   hSetBuffering stdout NoBuffering -- DO NOT REMOVE
+
+   flip evalStateT initState $ forever $ do
+      blocks <- liftIO (replicateM 8 getBlock)
+
+      myGrid <- liftIO getGrid
+      opGrid <- liftIO getGrid
+
+      ((col, rot), debug) <- runWriterT (step blocks myGrid)
+
+      liftIO $ mapM_ (hPutStrLn stderr) debug
+      liftIO $ putStrLn $ unwords $ map show [col, rot]
+
+getBlock :: IO Block
+getBlock = do
+   [colorA, colorB] <- map read . words <$> getLine
+   pure (colorA, colorB)
+
+getGrid :: IO Grid
+getGrid = fmap (A.array ((0,0),(5,11)) . concat) $
+      forM [0..11] $ \row -> do
+         line <- getLine
+         pure [ ((col, row), cell ch) | (col, ch) <- zip [0..] line ]
+   where
+      cell '.' = Empty
+      cell '0' = Skull
+      cell ch  = Color (read [ch])
+
+type StepState = ()
+initState = ()
+
+step :: (Applicative m, MonadState StepState m, MonadWriter [String] m)
+     => [Block] -> Grid -> m (Column, Rotation)
+step blocks myGrid = do
+   s <- get
+   let try c rot = ((c,rot),) . score s (tail blocks) <$> simulate myGrid (head blocks) c rot
+   let candidates = catMaybes $ try <$> [0..5] <*> [0..3]
+   let best = if null candidates then ((0,3),-1)
+              else maximumBy (compare `on` snd) candidates
+   pure (fst best)
+
+evalWriterT :: Monad m => WriterT w m a -> m a
+evalWriterT m = liftM fst (runWriterT m)
+
+score :: StepState -> [Block] -> (Grid, Int) -> Int
+score s blocks (grid, points) = flip evalState s $ evalWriterT $
+   let loop [] grid' points' =
+          let free      = length $ filter (== Empty) $ A.elems grid'
+              nonSkulls = length $ filter (/= Skull) $ A.elems grid'
+              levels    = length $ takeWhile emptyLevel [0..11]
+              emptyLevel r = all (\c -> grid'!(c,r) == Empty) [0..5]
+          in pure (points' + 100*nonSkulls + 10*levels)
+       loop blocks' grid' points' = do
+          (col, rot) <- step (take 1 blocks') grid'
+          let mgrid'' = simulate grid' (head blocks') col rot
+          case mgrid'' of
+             Nothing -> pure (-1000000)
+             Just (grid'', pts) -> loop (tail blocks') grid'' (points' + pts)
+   in loop (take 3 blocks) grid points
+
+simulate :: Grid -> Block -> Column -> Rotation -> Maybe (Grid, Int)
+simulate grid (colorA, colorB) col rot
+   | not (A.inRange (A.bounds grid) crA) ||
+     not (A.inRange (A.bounds grid) crB) ||
+     grid!crA /= Empty || grid!crB /= Empty = Nothing
+   | otherwise = Just . second getSum . runWriter $ simFall startGrid 1
+   where
+      (crA, crB) = case rot of
+         0 -> ((col,0), (col+1,0))
+         1 -> ((col,1), (col,  0))
+         2 -> ((col,0), (col-1,0))
+         3 -> ((col,0), (col,  1))
+      startGrid = grid // [ (crB, Color colorB), (crA, Color colorA) ]
+
+simFall :: (Applicative m, MonadWriter (Sum Int) m) => Grid -> Int -> m Grid
+simFall grid = simDisappear newGrid
+   where
+      packColumn c = zipWith (\r x -> ((c,r),x)) [11,10..0]
+                   $ (++ repeat Empty)
+                   $ filter (/= Empty)
+                   $ map (\r -> grid!(c,r)) [11,10..0]
+      newGrid = A.array ((0,0),(5,11)) $ concatMap packColumn [0..5]
+
+simDisappear :: (Applicative m, MonadWriter (Sum Int) m) => Grid -> Int -> m Grid
+simDisappear grid stage = case null erased of
+      True  -> pure grid
+      False -> do
+         tell . Sum $ 10 * blocksCleared * scale
+         simFall erasedGrid (stage + 1)
+   where
+      adjacent (c1,r1) (c2,r2) = (c1 == c2 && (r1 == r2 - 1 || r1 == r2 + 1)) ||
+                                 (r1 == r2 && (c1 == c2 - 1 || c1 == c2 + 1))
+      adjacentMatch (cr1, Color x1) (cr2, Color x2) =
+         x1 == x2 && adjacent cr1 cr2
+      colorCells = filter (isColor    . snd) $ A.assocs grid
+      skullCells = filter ((== Skull) . snd) $ A.assocs grid
+      groups = connectedGroups adjacentMatch (S.fromList colorCells)
+      largeGroups = filter ((>= 4) . S.size) groups
+      erasedColors = concatMap S.toList largeGroups
+      erasedSkulls = filter (\(cr,_) -> any (adjacent cr . fst) erasedColors) skullCells
+      erased = erasedColors ++ erasedSkulls
+      erasedGrid = grid // map (second (const Empty)) erased
+      blocksCleared = length erasedColors
+      chainPower = if stage < 2 then 0 else 8 * 2^(stage-2)
+      uniqueColors = length . nub $ map snd erasedColors
+      colorBonus = if uniqueColors < 2 then 0 else 2^(uniqueColors-1)
+      groupBonus = sum (map (perGroupBonus . S.size) largeGroups)
+      perGroupBonus n = if n >= 11 then 8 else n - 4
+      scale = max 1 $ min 999 $ chainPower + colorBonus + groupBonus
+
+isColor :: Cell -> Bool
+isColor Empty     = False
+isColor Skull     = False
+isColor (Color _) = True
+
+connectedGroups :: Ord a => (a -> a -> Bool) -> Set a -> [Set a]
+connectedGroups p rem = case S.minView rem of
+   Nothing        -> []
+   Just (x, rem') ->
+      let go fringe others = case S.minView fringe of
+             Nothing           -> (S.empty, others)
+             Just (y, fringe') -> case S.partition (p y) others of
+                (adj, notAdj) -> first (S.insert y) $
+                   go (S.union fringe' adj) notAdj
+          (conn, notConn) = go (S.singleton x) rem'
+      in conn : connectedGroups p notConn