{-# Language ImportQualifiedPost, BlockArguments #-}
{-|
Module      : Main
Description : Day 17 solution
Copyright   : (c) Eric Mertens, 2020
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2020/day/17>

This automata implementation uses a few optimizations:

It precomputes the neighborhood map that can be quickly
translated to any position without recomputing the map.
Unioning one of these for each live cell gives us a Map
where each element tells us how many neighbors that
cell has. Taking the union of these maps is faster than
inserting individual neighbor elements.

This solution only needs to consider elements that have
any neighbors at all, and these values are determined
by only considering the live elements from the previous
generation.

The solution avoids checking if a particular cell in the
previous generation was alive or not unless it's neighbors
is known to be @2@, as this is the only time it matters.

This solution works both with a flexible, n-dimensional
list coordinate representation and also with more
efficient unpacked tuples of integers. The list version
is about 2x slower than the unpacked tuples.

On my MacBook Pro, part 2 of this problem runs in 50ms.

-}
module Main (main) where

import Control.Monad (replicateM)
import Data.Map (Map)
import Data.Map.Strict qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set

import Advent (getInputMap, times)
import Advent.Coord (Coord(..))

-- |
-- >>> :main
-- 257
-- 2532
main :: IO ()
IO ()
main =
  do Set Coord
inp <- Map Coord Char -> Set Coord
forall k a. Map k a -> Set k
Map.keysSet (Map Coord Char -> Set Coord)
-> (Map Coord Char -> Map Coord Char)
-> Map Coord Char
-> Set Coord
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Map Coord Char -> Map Coord Char
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Char
'#'Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==) (Map Coord Char -> Set Coord)
-> IO (Map Coord Char) -> IO (Set Coord)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO (Map Coord Char)
getInputMap Int
2020 Int
17
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((C3 -> Map C3 Int) -> Set C3 -> Int
forall a. Ord a => (a -> Map a Int) -> Set a -> Int
run C3 -> Map C3 Int
neighborCount3 ((Coord -> C3) -> Set Coord -> Set C3
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map Coord -> C3
toC3 Set Coord
inp))
     Int -> IO ()
forall a. Show a => a -> IO ()
print ((C4 -> Map C4 Int) -> Set C4 -> Int
forall a. Ord a => (a -> Map a Int) -> Set a -> Int
run C4 -> Map C4 Int
neighborCount4 ((Coord -> C4) -> Set Coord -> Set C4
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map Coord -> C4
toC4 Set Coord
inp))

run ::
  Ord a =>
  (a -> Map a Int) {- ^ neighbor generator             -} ->
  Set a            {- ^ input coordinates              -} ->
  Int              {- ^ live cells after 6 generations -}
run :: forall a. Ord a => (a -> Map a Int) -> Set a -> Int
run a -> Map a Int
neighbor = Set a -> Int
forall a. Set a -> Int
Set.size (Set a -> Int) -> (Set a -> Set a) -> Set a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Set a -> Set a) -> Set a -> Set a
forall a. Int -> (a -> a) -> a -> a
times Int
6 ((a -> Map a Int) -> Set a -> Set a
forall a. Ord a => (a -> Map a Int) -> Set a -> Set a
step a -> Map a Int
neighbor)

-- | Determine if a cell should be alive in the next generation.
rule ::
  Ord a =>
  Set a {- ^ previous generation      -} ->
  a     {- ^ coordinate               -} ->
  Int   {- ^ live neighbor count      -} ->
  Bool  {- ^ alive in next generation -}
rule :: forall a. Ord a => Set a -> a -> Int -> Bool
rule Set a
world a
c Int
n = Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 Bool -> Bool -> Bool
&& a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member a
c Set a
world

-- | Compute the next generation from the previous generation
step :: Ord a => (a -> Map a Int) -> Set a -> Set a
step :: forall a. Ord a => (a -> Map a Int) -> Set a -> Set a
step a -> Map a Int
neighbor Set a
world
  = Map a Int -> Set a
forall k a. Map k a -> Set k
Map.keysSet
  (Map a Int -> Set a) -> Map a Int -> Set a
forall a b. (a -> b) -> a -> b
$ (a -> Int -> Bool) -> Map a Int -> Map a Int
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (Set a -> a -> Int -> Bool
forall a. Ord a => Set a -> a -> Int -> Bool
rule Set a
world)
  (Map a Int -> Map a Int) -> Map a Int -> Map a Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> [Map a Int] -> Map a Int
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)
  ([Map a Int] -> Map a Int) -> [Map a Int] -> Map a Int
forall a b. (a -> b) -> a -> b
$ (a -> Map a Int) -> [a] -> [Map a Int]
forall a b. (a -> b) -> [a] -> [b]
map a -> Map a Int
neighbor
  ([a] -> [Map a Int]) -> [a] -> [Map a Int]
forall a b. (a -> b) -> a -> b
$ Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
world

-- List-based coordinates ----------------------------------------------

-- Given a dimension and a coordinate this produces a map of all the
-- neighboring cells with values of @1@.
--
-- neighborCountN :: Int -> [Int] -> Map [Int] Int
-- neighborCountN d =
--   let ns = Map.fromList [(c,1) | c <- tail (replicateM d [0,-1,1])]
--   in \c -> Map.mapKeysMonotonic (zipWith (+) c) ns

-- Unpacked 3-tuples ---------------------------------------------------

data C3 = C3 !Int !Int !Int      deriving (C3 -> C3 -> Bool
(C3 -> C3 -> Bool) -> (C3 -> C3 -> Bool) -> Eq C3
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: C3 -> C3 -> Bool
== :: C3 -> C3 -> Bool
$c/= :: C3 -> C3 -> Bool
/= :: C3 -> C3 -> Bool
Eq, Eq C3
Eq C3 =>
(C3 -> C3 -> Ordering)
-> (C3 -> C3 -> Bool)
-> (C3 -> C3 -> Bool)
-> (C3 -> C3 -> Bool)
-> (C3 -> C3 -> Bool)
-> (C3 -> C3 -> C3)
-> (C3 -> C3 -> C3)
-> Ord C3
C3 -> C3 -> Bool
C3 -> C3 -> Ordering
C3 -> C3 -> C3
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: C3 -> C3 -> Ordering
compare :: C3 -> C3 -> Ordering
$c< :: C3 -> C3 -> Bool
< :: C3 -> C3 -> Bool
$c<= :: C3 -> C3 -> Bool
<= :: C3 -> C3 -> Bool
$c> :: C3 -> C3 -> Bool
> :: C3 -> C3 -> Bool
$c>= :: C3 -> C3 -> Bool
>= :: C3 -> C3 -> Bool
$cmax :: C3 -> C3 -> C3
max :: C3 -> C3 -> C3
$cmin :: C3 -> C3 -> C3
min :: C3 -> C3 -> C3
Ord)

toC3 :: Coord -> C3
toC3 :: Coord -> C3
toC3 (C Int
y Int
x) = Int -> Int -> Int -> C3
C3 Int
x Int
y Int
0

-- | Compute a Map with @1@ stored at each neighboring coordinate
neighborCount3 :: C3 -> Map C3 Int
neighborCount3 :: C3 -> Map C3 Int
neighborCount3 =
  let ns :: Map C3 Int
ns = [(C3, Int)] -> Map C3 Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Int -> Int -> Int -> C3
C3 Int
x Int
y Int
z,Int
1) | [Int
x,Int
y,Int
z] <- [[Int]] -> [[Int]]
forall a. HasCallStack => [a] -> [a]
tail (Int -> [Int] -> [[Int]]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
3 [Int
0,-Int
1,Int
1])]
  in \(C3 Int
a Int
b Int
c) -> (C3 -> C3) -> Map C3 Int -> Map C3 Int
forall k1 k2 a. (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeysMonotonic (\(C3 Int
x Int
y Int
z) -> Int -> Int -> Int -> C3
C3 (Int
aInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
x) (Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
y) (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
z)) Map C3 Int
ns

-- Unpacked 4-tuples ---------------------------------------------------

data C4 = C4 !Int !Int !Int !Int deriving (C4 -> C4 -> Bool
(C4 -> C4 -> Bool) -> (C4 -> C4 -> Bool) -> Eq C4
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: C4 -> C4 -> Bool
== :: C4 -> C4 -> Bool
$c/= :: C4 -> C4 -> Bool
/= :: C4 -> C4 -> Bool
Eq, Eq C4
Eq C4 =>
(C4 -> C4 -> Ordering)
-> (C4 -> C4 -> Bool)
-> (C4 -> C4 -> Bool)
-> (C4 -> C4 -> Bool)
-> (C4 -> C4 -> Bool)
-> (C4 -> C4 -> C4)
-> (C4 -> C4 -> C4)
-> Ord C4
C4 -> C4 -> Bool
C4 -> C4 -> Ordering
C4 -> C4 -> C4
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: C4 -> C4 -> Ordering
compare :: C4 -> C4 -> Ordering
$c< :: C4 -> C4 -> Bool
< :: C4 -> C4 -> Bool
$c<= :: C4 -> C4 -> Bool
<= :: C4 -> C4 -> Bool
$c> :: C4 -> C4 -> Bool
> :: C4 -> C4 -> Bool
$c>= :: C4 -> C4 -> Bool
>= :: C4 -> C4 -> Bool
$cmax :: C4 -> C4 -> C4
max :: C4 -> C4 -> C4
$cmin :: C4 -> C4 -> C4
min :: C4 -> C4 -> C4
Ord)

toC4 :: Coord -> C4
toC4 :: Coord -> C4
toC4 (C Int
y Int
x) = Int -> Int -> Int -> Int -> C4
C4 Int
x Int
y Int
0 Int
0

neighborCount4 :: C4 -> Map C4 Int
neighborCount4 :: C4 -> Map C4 Int
neighborCount4 =
  let ns :: Map C4 Int
ns = [(C4, Int)] -> Map C4 Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Int -> Int -> Int -> Int -> C4
C4 Int
x Int
y Int
z Int
w,Int
1) | [Int
x,Int
y,Int
z,Int
w] <- [[Int]] -> [[Int]]
forall a. HasCallStack => [a] -> [a]
tail (Int -> [Int] -> [[Int]]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
4 [Int
0,-Int
1,Int
1])]
  in \(C4 Int
a Int
b Int
c Int
d) -> (C4 -> C4) -> Map C4 Int -> Map C4 Int
forall k1 k2 a. (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeysMonotonic (\(C4 Int
x Int
y Int
z Int
w) -> Int -> Int -> Int -> Int -> C4
C4 (Int
aInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
x) (Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
y) (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
z) (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
w)) Map C4 Int
ns