day 15: rewrite part 2

Signed-off-by: Amneesh Singh <natto@weirdnatto.in>

day14.hs

@@ -3,41 +3,40 @@
 
 import Data.Maybe (fromJust)
 import Data.Set (Set)
-import qualified Data.Set as S (fromList, insert, lookupMax, member, notMember, union)
+import qualified Data.Set as S (fromList, insert, map, member, notMember, union)
 import Data.Text (Text)
 import qualified Data.Text as T (lines, splitOn, words)
 import qualified Data.Text.Read as T (decimal)
 import Lib (readFile')
 
-
--- ok this is it, this takes 4.54s, kms
--- this can probably be improved using hashtables in ST monad but will try that later
+-- times can be improved greatly using Data.HashSet in unordered-containers, but i wont use it
 
 main :: IO ()
 main = do
   input <- parse . T.lines <$> readFile' "day14.in"
-  let mx = fst . fromJust $ S.lookupMax input
+  let mx = maximum $ S.map snd input
   putStr "Q1: "
-  print $ sand (0, 500) input 0 mx Q1
+  print $ sand (500, 0) input 0 mx Q1
   putStr "Q2: "
-  print $ sand (0, 500) input 0 (mx + 2) Q2
+  print $ sand (500, 0) input 0 (mx + 2) Q2
+
+type Coord = (Int, Int)
 
 data Q = Q1 | Q2 deriving (Eq)
 
-sand :: (Int, Int) -> Set (Int, Int) -> Int -> Int -> Q -> Int
-sand (y, x) rocks soFar mx q
-  | (q == Q1 && y > mx) || S.member (y, x) rocks = soFar
-  | q == Q2 && y + 1 == mx = sand (0, 500) (S.insert (y, x) rocks) (soFar + 1) mx q
-  | S.notMember (y + 1, x) rocks = sand (y + 1, x) rocks soFar mx q
-  | S.notMember (y + 1, x - 1) rocks = sand (y + 1, x - 1) rocks soFar mx q
-  | S.notMember (y + 1, x + 1) rocks = sand (y + 1, x + 1) rocks soFar mx q
-  | otherwise = sand (0, 500) (S.insert (y, x) rocks) (soFar + 1) mx q
+sand :: Coord -> Set Coord -> Int -> Int -> Q -> Int
+sand (x, y) rocks soFar mx q
+  | (q == Q1 && y > mx) || S.member (x, y) rocks = soFar
+  | q == Q2 && y + 1 == mx = sand (500, 0) (S.insert (x, y) rocks) (soFar + 1) mx q
+  | S.notMember (x, y + 1) rocks = sand (x, y + 1) rocks soFar mx q
+  | S.notMember (x - 1, y + 1) rocks = sand (x - 1, y + 1) rocks soFar mx q
+  | S.notMember (x + 1, y + 1) rocks = sand (x + 1, y + 1) rocks soFar mx q
+  | otherwise = sand (500, 0) (S.insert (x, y) rocks) (soFar + 1) mx q
 
-parse :: [Text] -> Set (Int, Int)
+parse :: [Text] -> Set Coord
 parse =
   S.fromList
-    . concat
-    . map
+    . concatMap
       ( ranges
           . map
             ( (\[x, y] -> (x, y))
@@ -48,14 +47,13 @@ parse =
           . T.words
       )
   where
-    ranges :: [(Int, Int)] -> [(Int, Int)]
+    ranges :: [Coord] -> [Coord]
     ranges xs =
-      concat $
-        foldr
-          ( \((a, b), (x, y)) z ->
-              if a == x
-                then zip [min b y .. max b y] (repeat a) : z
-                else zip (repeat b) [min a x .. max a x] : z
-          )
-          []
-          (zip xs (tail xs))
+      foldr
+        ( \((a, b), (x, y)) z ->
+            if a == x
+              then zip (repeat a) [min b y .. max b y] ++ z
+              else zip [min a x .. max a x] (repeat b) ++ z
+        )
+        []
+        (zip xs (tail xs))

day15.hs

@@ -0,0 +1,115 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ViewPatterns #-}
+
+import Data.List (nub, sort)
+import Data.Maybe (catMaybes)
+import Data.Text (Text)
+import qualified Data.Text as T (lines, null, split)
+import qualified Data.Text.Read as T (decimal, signed)
+import Lib (readFile')
+
+main :: IO ()
+main = do
+  input <- map parse . T.lines <$> readFile' "day15.in"
+  putStr "Q1: "
+  print $ q1 input 2000000
+  putStr "Q2: "
+  print $ q2 input 0 4000000
+
+type Coords = (Int, Int, Int, Int, Int)
+
+q1 :: [Coords] -> Int -> Int
+q1 input n = mergeRanges getRanges
+  where
+    mergeRanges :: [(Int, Int)] -> Int
+    mergeRanges [(x0, x1)] = x1 - x0 + 1
+    mergeRanges ((x0, x1) : (x2, x3) : xs)
+      | x2 - x1 < 2 = mergeRanges $ (x0, max x1 x3) : xs
+      | otherwise = x1 - x0 + 1 + mergeRanges ((x2, x3) : xs)
+
+    getRanges :: [(Int, Int)]
+    getRanges =
+      sort $
+        foldr
+          ( \(x, y, r, a, b) xs -> case r - abs (n - y) of
+              dx
+                | dx < 0 -> xs
+                | n == b -> (x - dx, a - 1) : (a + 1, x + dx) : xs
+                | otherwise -> (x - dx, x + dx) : xs
+          )
+          []
+          input
+
+type Segment = (Int, Int, Int, Int)
+
+q2 input mn mx =
+  head
+    [ a * mx + b
+      | (a, b) <- diag1234,
+        a >= mn && b >= mn && a <= mx && b <= mx,
+        and [abs (b - y) + abs (a - x) > r | (x, y, r, _, _) <- input]
+    ]
+  where
+    diag1, diag2, diag3, diag4 :: Coords -> Segment
+    diag1 (x, y, r, _, _) = (x - r - 1, y, x, y - r - 1)
+    diag2 (x, y, r, _, _) = (x, y - r - 1, x + r + 1, y)
+    diag3 (x, y, r, _, _) = (x, y + r + 1, x + r + 1, y)
+    diag4 (x, y, r, _, _) = (x - r - 1, y, x, y + r + 1)
+
+    diagInt :: Segment -> Segment -> Maybe Segment
+    diagInt a@(x0, y0, x1, y1) b@(x2, y2, x3, y3)
+      | x3 < x0 || x1 < x2 = Nothing
+      | x0 >= x2 && x1 <= x3 = Just a
+      | x2 >= x0 && x3 <= x1 = Just b
+      | x3 >= x1 = Just (x2, y2, x1, y1)
+      | x1 >= x3 = Just (x0, y0, x3, y3)
+
+    diag13, diag24 :: [Segment]
+    diag13 =
+      catMaybes
+        [ diagInt i j
+          | a <- input,
+            b <- input,
+            let i@(x0, y0, x1, y1) = diag1 a,
+            let j@(x2, y2, x3, y3) = diag3 b,
+            a /= b,
+            x0 + y0 == x2 + y2
+        ]
+    diag24 =
+      catMaybes
+        [ diagInt i j
+          | a <- input,
+            b <- input,
+            let i@(x0, y0, x1, y1) = diag2 a,
+            let j@(x2, y2, x3, y3) = diag4 b,
+            a /= b,
+            y0 - x0 == y2 - x2
+        ]
+
+    diag1234 :: [(Int, Int)]
+    diag1234 =
+      [ let (c1, c2) = (y0 + x0, y2 - x2) in (div (c1 - c2) 2, div (c1 + c2) 2)
+        | (x0, y0, x1, y1) <- diag13,
+          (x2, y2, x3, y3) <- diag24
+      ]
+
+parse :: Text -> Coords
+parse =
+  ( \[ "Sensor",
+       "at",
+       "x",
+       T.signed T.decimal -> Right (x, ""),
+       "y",
+       T.signed T.decimal -> Right (y, ""),
+       "closest",
+       "beacon",
+       "is",
+       "at",
+       "x",
+       T.signed T.decimal -> Right (a, ""),
+       "y",
+       T.signed T.decimal -> Right (b, "")
+       ] -> (x, y, abs (b - y) + abs (a - x), a, b)
+  )
+    . filter (not . T.null)
+    . T.split (`elem` (":=, " :: [Char]))