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

<https://adventofcode.com/2022/day/9>

>>> :{
:main +
  "R 4\n\
  \U 4\n\
  \L 3\n\
  \D 1\n\
  \R 4\n\
  \D 1\n\
  \L 5\n\
  \R 2\n"
:}
13
1

>>> :{
:main +
  "R 5\n\
  \U 8\n\
  \L 8\n\
  \D 3\n\
  \R 17\n\
  \D 10\n\
  \L 25\n\
  \U 20\n"
:}
88
36

>>> uniqueLocations [(CR,4),(CU,4),(CL,3),(CD,1),(CR,4),(CD,1),(CL,5),(CR,2)] !! 1
13

>>> uniqueLocations [(CR,5),(CU,8),(CL,8),(CD,3),(CR,17),(CD,10),(CL,25),(CU,20)] !! 9
36

-}
module Main where

import Data.List (transpose)
import Data.Set qualified as Set

import Advent (format, stageTH)
import Advent.Coord (Coord(..), origin, east, north, south, west, normInf)

-- | Rope movement instructions
data C
  = CD -- ^ move down
  | CR -- ^ move right
  | CU -- ^ move up
  | CL -- ^ move left
  deriving Int -> C -> ShowS
[C] -> ShowS
C -> String
(Int -> C -> ShowS) -> (C -> String) -> ([C] -> ShowS) -> Show C
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> C -> ShowS
showsPrec :: Int -> C -> ShowS
$cshow :: C -> String
show :: C -> String
$cshowList :: [C] -> ShowS
showList :: [C] -> ShowS
Show

stageTH

-- | Print the answers to both parts of day 9. Automatically finds
-- input file unless overridden with a command line argument.
--
-- >>> :main
-- 5930
-- 2443
main :: IO ()
IO ()
main =
 do [(C, Int)]
input <- [format|2022 9 (@C %u%n)*|]
    let knots :: [Int]
knots = [(C, Int)] -> [Int]
uniqueLocations [(C, Int)]
input
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int]
knots [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1)
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int]
knots [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
9)

-- | Generate the number of unique locations each knot in an infinitely long
-- rope visits given a list of movement commands.
uniqueLocations :: [(C,Int)] -> [Int]
uniqueLocations :: [(C, Int)] -> [Int]
uniqueLocations
  = ([Coord] -> Int) -> [[Coord]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Coord] -> Int
forall a. Ord a => [a] -> Int
countUnique                -- list of unique locations per knot
  ([[Coord]] -> [Int])
-> ([(C, Int)] -> [[Coord]]) -> [(C, Int)] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Coord]] -> [[Coord]]
forall a. [[a]] -> [[a]]
transpose                      -- list of steps to list of knots
  ([[Coord]] -> [[Coord]])
-> ([(C, Int)] -> [[Coord]]) -> [(C, Int)] -> [[Coord]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Coord] -> C -> [Coord]) -> [Coord] -> [C] -> [[Coord]]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl [Coord] -> C -> [Coord]
stepRope (Coord -> [Coord]
forall a. a -> [a]
repeat Coord
origin) -- step a rope starting at origin through list of movements
  ([C] -> [[Coord]])
-> ([(C, Int)] -> [C]) -> [(C, Int)] -> [[Coord]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((C, Int) -> [C]) -> [(C, Int)] -> [C]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(C
c,Int
n) -> Int -> C -> [C]
forall a. Int -> a -> [a]
replicate Int
n C
c)

-- | Generate the unit vector corresponding to an input command.
cToVec :: C -> Coord
cToVec :: C -> Coord
cToVec C
CU = Coord
north
cToVec C
CD = Coord
south
cToVec C
CR = Coord
east
cToVec C
CL = Coord
west

-- | Update all the knot locations in a rope given a step direction for the head knot.
stepRope ::
  [Coord] {- ^ knot locations         -} ->
  C       {- ^ next step direction    -} ->
  [Coord] {- ^ updated knot locations -}
stepRope :: [Coord] -> C -> [Coord]
stepRope (Coord
x:[Coord]
xs) C
c = Coord -> [Coord] -> [Coord]
updateTails (C -> Coord
cToVec C
c Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+ Coord
x) [Coord]
xs
stepRope []     C
_ = []

-- | Update all the tail knots in the rope given a new head position.
updateTails ::
  Coord   {- ^ head         -} ->
  [Coord] {- ^ tails        -} ->
  [Coord] {- ^ updated rope -}
updateTails :: Coord -> [Coord] -> [Coord]
updateTails Coord
h [] = [Coord
h]
updateTails Coord
h (Coord
t : [Coord]
ts)
  | Coord -> Bool
isNearOrigin Coord
delta = Coord
h Coord -> [Coord] -> [Coord]
forall a. a -> [a] -> [a]
: Coord
t Coord -> [Coord] -> [Coord]
forall a. a -> [a] -> [a]
: [Coord]
ts -- once a knot is stationary, the rest will be, too
  | Bool
otherwise          = Coord
h Coord -> [Coord] -> [Coord]
forall a. a -> [a] -> [a]
: Coord -> [Coord] -> [Coord]
updateTails (Coord
t Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+ Coord -> Coord
forall a. Num a => a -> a
signum Coord
delta) [Coord]
ts
  where
    delta :: Coord
delta = Coord
h Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
- Coord
t

-- | Predicate for coordinates at or adjacent to the origin.
--
-- >>> all isNearOrigin [C y x | y <- [-1..1], x <- [-1..1]]
-- True
--
-- >>> any isNearOrigin [C 2 0, C 0 2, C 2 1, C (-2) 0, C (-1) 2]
-- False
isNearOrigin :: Coord -> Bool
isNearOrigin :: Coord -> Bool
isNearOrigin Coord
c = Coord -> Int
normInf Coord
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2

-- | Return the number of unique elements in a list.
countUnique :: Ord a => [a] -> Int
countUnique :: forall a. Ord a => [a] -> Int
countUnique = Set a -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Set a -> Int) -> ([a] -> Set a) -> [a] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList