Smash The Code Contest -- Silver league, rank 5

SmashTheCode/SmashTheCode.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE FlexibleContexts, LambdaCase, TupleSections #-}
+{-# LANGUAGE FlexibleContexts, LambdaCase, TupleSections, UnicodeSyntax #-}
 
 import Control.Applicative
 import Control.Arrow (first, second, (&&&))
@@ -12,6 +12,7 @@ import Data.Monoid
 import Data.Array (Array, (!), (//))
 import Data.Set (Set)
 import System.IO
+import System.Timeout
 
 import qualified Data.Array       as A
 import qualified Data.Foldable    as F
@@ -35,26 +36,26 @@ main = do
    hSetBuffering stdout NoBuffering -- DO NOT REMOVE
 
    flip evalStateT initState $ forever $ do
-      blocks <- liftIO (replicateM 8 getBlock)
+      blocks ← liftIO (replicateM 8 getBlock)
 
-      myGrid <- liftIO getGrid
+      myGrid ← liftIO getGrid
-      opGrid <- liftIO getGrid
+      opGrid ← liftIO getGrid
 
-      ((col, rot), debug) <- runWriterT (step blocks myGrid)
+      ((col, rot), debug) ← runWriterT (step blocks myGrid)
 
-      liftIO $ mapM_ (hPutStrLn stderr) debug
+      --liftIO $ mapM_ (hPutStrLn stderr) debug
       liftIO $ putStrLn $ unwords $ map show [col, rot]
 
 getBlock :: IO Block
 getBlock = do
-   [colorA, colorB] <- map read . words <$> getLine
+   [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
+      forM [0..11] $ \row → do
-         line <- getLine
+         line ← getLine
-         pure [ ((col, row), cell ch) | (col, ch) <- zip [0..] line ]
+         pure [ ((col, row), cell ch) | (col, ch) ← zip [0..] line ]
    where
       cell '.' = Empty
       cell '0' = Skull
@@ -64,35 +65,49 @@ type StepState = ()
 initState = ()
 
 step :: (Applicative m, MonadState StepState m, MonadWriter [String] m)
-     => [Block] -> Grid -> m (Column, Rotation)
+     ⇒ [Block] → Grid → m (Column, Rotation)
-step blocks myGrid = do
+step blocks myGrid = fst <$> step' 0 blocks myGrid
-   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
+step' :: (Applicative m, MonadState StepState m, MonadWriter [String] m)
-evalWriterT m = liftM fst (runWriterT m)
+      ⇒ Int → [Block] → Grid → m ((Column, Rotation), Int)
+step' depth (block:blocks) myGrid = do
+   let try grid bl c rot = do
+         result ← simulate grid bl c rot
+         pure (score result, (result, (c, rot)))
 
-score :: StepState -> [Block] -> (Grid, Int) -> Int
+   let candidates = catMaybes $ try myGrid block <$> [0..5] <*> [0..3]
-score s blocks (grid, points) = flip evalState s $ evalWriterT $
+   let best = sortBy (flip compare `on` fst) candidates
-   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)
+   let limit = [5,2,1,1,1,1,0] !! depth
+   best' ← if length (take limit best) < 1 || null blocks then pure best else do
+      s ← get
+      candidates' ← forM (take limit best) $
+         \(score1, ((grid', points), (c, rot))) → do
+            let ((_, score2), w) = flip evalState s $ runWriterT $
+                                   step' (depth + 1) blocks grid'
+            tell w
+            pure (score1 + score2, ((grid', points), (c, rot)))
+      pure $ sortBy (flip compare `on` fst) candidates'
+
+   -- tell [show depth ++ ": " ++ show (map fst $ take limit best')]
+
+   case best' of
+      [] → pure ((0, 0), -1000000)
+      ((score1, (_, move1)):_) → pure (move1, score1)
+
+score :: (Grid, Int) → Int
+score (grid, points) = 1000*points + 5*nonSkulls + sum groups
+   where
+      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]
+      colorCells = filter (isColor . snd) $ A.assocs grid
+      groups = map (\g -> (S.size g - 1)^2)
+             $ connectedGroups adjacentMatch
+             $ S.fromList colorCells
+
+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) ||
@@ -100,38 +115,34 @@ simulate grid (colorA, colorB) col rot
    | otherwise = Just . second getSum . runWriter $ simFall startGrid 1
    where
       (crA, crB) = case rot of
-         0 -> ((col,0), (col+1,0))
+         0 → ((col,0), (col+1,0))
-         1 -> ((col,1), (col,  0))
+         1 → ((col,1), (col,  0))
-         2 -> ((col,0), (col-1,0))
+         2 → ((col,0), (col-1,0))
-         3 -> ((col,0), (col,  1))
+         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 :: (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]
+      packColumn c = zipWith (\r x → ((c,r),x)) [11,10..0]
                    $ (++ repeat Empty)
                    $ filter (/= Empty)
-                   $ map (\r -> grid!(c,r)) [11,10..0]
+                   $ 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 :: (Applicative m, MonadWriter (Sum Int) m) ⇒ Grid → Int → m Grid
 simDisappear grid stage = case null erased of
-      True  -> pure grid
+      True  → pure grid
-      False -> do
+      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
+      erasedSkulls = filter (\(cr,_) → any (adjacent cr . fst) erasedColors) skullCells
       erased = erasedColors ++ erasedSkulls
       erasedGrid = grid // map (second (const Empty)) erased
       blocksCleared = length erasedColors
@@ -142,19 +153,26 @@ simDisappear grid stage = case null erased of
       perGroupBonus n = if n >= 11 then 8 else n - 4
       scale = max 1 $ min 999 $ chainPower + colorBonus + groupBonus
 
-isColor :: Cell -> Bool
+isColor :: Cell → Bool
 isColor Empty     = False
 isColor Skull     = False
 isColor (Color _) = True
 
-connectedGroups :: Ord a => (a -> a -> Bool) -> Set a -> [Set a]
+adjacent :: (Column,Row) → (Column,Row) → Bool
+adjacent (c1,r1) (c2,r2) = (c1 == c2 && (r1 == r2 - 1 || r1 == r2 + 1)) ||
+                           (r1 == r2 && (c1 == c2 - 1 || c1 == c2 + 1))
+
+adjacentMatch :: ((Column, Row), Cell) → ((Column, Row), Cell) → Bool
+adjacentMatch (cr1,x1) (cr2,x2) = x1 == x2 && adjacent cr1 cr2
+
+connectedGroups :: Ord a ⇒ (a → a → Bool) → Set a → [Set a]
 connectedGroups p rem = case S.minView rem of
-   Nothing        -> []
+   Nothing        → []
-   Just (x, rem') ->
+   Just (x, rem') →
       let go fringe others = case S.minView fringe of
-             Nothing           -> (S.empty, others)
+             Nothing           → (S.empty, others)
-             Just (y, fringe') -> case S.partition (p y) others of
+             Just (y, fringe') → case S.partition (p y) others of
-                (adj, notAdj) -> first (S.insert y) $
+                (adj, notAdj) → first (S.insert y) $
                    go (S.union fringe' adj) notAdj
           (conn, notConn) = go (S.singleton x) rem'
       in conn : connectedGroups p notConn