CodinGame/SmashTheCode/SmashTheCode.hs

179 lines
6.6 KiB
Haskell

{-# LANGUAGE FlexibleContexts, LambdaCase, TupleSections, UnicodeSyntax #-}
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 System.Timeout
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 = fst <$> step' 0 blocks myGrid
step' :: (Applicative m, MonadState StepState m, MonadWriter [String] 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)))
let candidates = catMaybes $ try myGrid block <$> [0..5] <*> [0..3]
let best = sortBy (flip compare `on` fst) candidates
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) ||
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
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
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 []
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