149 lines
3.7 KiB
Haskell
149 lines
3.7 KiB
Haskell
{-# options_ghc -Wall #-}
|
||
{-# language BlockArguments, LambdaCase #-}
|
||
import Data.List
|
||
import System.Environment
|
||
import qualified Data.Map as M
|
||
import Data.Maybe
|
||
import qualified Data.Array as A
|
||
|
||
|
||
data Block = Rock | Space | Wall deriving (Eq, Ord, Show)
|
||
|
||
|
||
data Rotation = U | L | D | R deriving (Eq, Ord, Show, Bounded, Enum)
|
||
|
||
left, right, opp :: Rotation -> Rotation
|
||
left = \case U -> L; L -> D; D -> R; R -> U
|
||
right = \case U -> R; R -> D; D -> L; L -> U
|
||
opp = \case U -> D; D -> U; L -> R; R -> L
|
||
|
||
instance Semigroup Rotation where
|
||
U <> r = r
|
||
L <> r = left r
|
||
R <> r = right r
|
||
D <> r = opp r
|
||
|
||
instance Monoid Rotation where mempty = U
|
||
|
||
type I2 = (Int, Int)
|
||
|
||
data RArray e = RA !Rotation !(A.Array I2 e)
|
||
|
||
rotateI :: Rotation -> (I2, I2) -> I2 -> I2
|
||
rotateI U _ i = i
|
||
rotateI L (_, (xhi, _)) (x, y) = (y, xhi - x - 1)
|
||
rotateI D (_, (xhi, yhi)) (x, y) = (xhi - x - 1, yhi - y - 1)
|
||
rotateI R (_, (_, yhi)) (x, y) = (yhi - y - 1, x)
|
||
|
||
(!) :: RArray e -> I2 -> e
|
||
RA r arr ! i = arr A.! rotateI r (A.bounds arr) i
|
||
|
||
rotate :: Rotation -> RArray e -> RArray e
|
||
rotate r1 (RA r2 arr) = RA (r1 <> r2) arr
|
||
|
||
|
||
block :: Char -> Maybe Block
|
||
block 'O' = Just Rock
|
||
block '.' = Just Space
|
||
block '#' = Just Wall
|
||
block _ = Nothing
|
||
|
||
type Board = [[Block]]
|
||
|
||
board :: String -> Board
|
||
board = map (map block') . lines where
|
||
block' :: Char -> Block
|
||
block' c = fromMaybe (error $ "unrecognised " ++ show c) $ block c
|
||
|
||
chunks :: (a -> Bool) -> [a] -> [[a]]
|
||
chunks _ [] = []
|
||
chunks want (x : xs)
|
||
| want x = let (as, bs) = span want xs in (x : as) : chunks want bs
|
||
| otherwise = let (as, bs) = break want xs in (x : as) : chunks want bs
|
||
|
||
fall :: Ord a => (a -> Bool) -> [a] -> [a]
|
||
fall p = concat . map sort . chunks p
|
||
|
||
rotateLeft :: [[a]] -> [[a]]
|
||
rotateLeft = transpose . map reverse
|
||
|
||
rotateRight :: [[a]] -> [[a]]
|
||
rotateRight = map reverse . transpose
|
||
|
||
rotate180 :: [[a]] -> [[a]]
|
||
rotate180 = reverse . map reverse
|
||
|
||
fallLeft, fallUp, fallRight, fallDown :: Board -> Board
|
||
fallLeft = map $ fall (/= Wall)
|
||
fallUp = rotateRight . fallLeft . rotateLeft
|
||
fallRight = rotate180 . fallLeft . rotate180
|
||
fallDown = rotateLeft . fallLeft . rotateRight
|
||
|
||
scoreWith :: (a -> Bool) -> [[a]] -> Int
|
||
scoreWith p xss =
|
||
let ℓ = length xss; ℓs = [ℓ, ℓ-1 .. 0] in
|
||
sum $ zipWith (\xs y -> length (filter p xs) * y) xss ℓs
|
||
|
||
score :: Board -> Int
|
||
score = scoreWith (== Rock)
|
||
|
||
part1 :: Board -> Int
|
||
part1 = score . fallUp
|
||
|
||
type Cycle = (Board, Board, Board, Board)
|
||
|
||
cycle1 :: Board -> Cycle
|
||
cycle1 xs =
|
||
let as = fallUp xs
|
||
bs = fallLeft as
|
||
cs = fallDown bs
|
||
ds = fallRight cs in
|
||
(as, bs, cs, ds)
|
||
|
||
nextCycle :: Cycle -> Cycle
|
||
nextCycle (_, _, _, xs) = cycle1 xs
|
||
|
||
cycles :: Board -> [Cycle]
|
||
cycles start = let c = cycle1 start in iterate nextCycle c
|
||
|
||
fourth :: (a,b,c,d) -> d
|
||
fourth (_,_,_,d) = d
|
||
|
||
findRepeat :: Ord a =>
|
||
M.Map a Int -> M.Map Int a ->
|
||
Int -> [a] -> (Int, Int, M.Map Int a)
|
||
findRepeat ai ia i (a:as)
|
||
| Just j <- M.lookup a ai = (j, i, ia)
|
||
| otherwise = findRepeat ai' ia' (i + 1) as
|
||
where ai' = M.insert a i ai; ia' = M.insert i a ia
|
||
findRepeat _ _ _ _ = error "you said it was infinite"
|
||
|
||
part2 :: Board -> Int
|
||
part2 ℓs =
|
||
let loops = 1_000_000_000
|
||
(m, n, xs) = findRepeat M.empty M.empty 0 $ cycles ℓs
|
||
index = (loops - n) `mod` (n - m) - 1 in
|
||
score $ fourth $ xs M.! (m + index)
|
||
|
||
main :: IO ()
|
||
main = do
|
||
[part, file] <- getArgs
|
||
ℓs <- board <$> readFile file
|
||
case part of
|
||
"1" -> print $ part1 ℓs
|
||
"2" -> print $ part2 ℓs
|
||
_ -> error "usage: $0 <part> <file>"
|
||
|
||
-- example1 :: [String]
|
||
-- example1 =
|
||
-- ["O....#....",
|
||
-- "O.OO#....#",
|
||
-- ".....##...",
|
||
-- "OO.#O....O",
|
||
-- ".O.....O#.",
|
||
-- "O.#..O.#.#",
|
||
-- "..O..#O..O",
|
||
-- ".......O..",
|
||
-- "#....###..",
|
||
-- "#OO..#...."]
|