Day twelve part two

day-12-haskell/app/Main.hs

@@ -1,7 +1,7 @@
 module Main (main) where
 
 import System.Environment.Blank ( getArgs )
-import Lib (task1)
+import Lib (task1, task2)
 
 main :: IO ()
 main = do
@@ -11,7 +11,8 @@ main = do
       input <- readFile file
       result <- task1 input
       print result
-    -- ["2", file] -> do
-    --   input <- readFile file
-    --   print $ task2 input
+    ["2", file] -> do
+      input <- readFile file
+      result <- task2 input
+      print result
     _ -> error "Usage: <1|2> <input file>"

day-12-haskell/src/Lib.hs

@@ -1,30 +1,34 @@
 {-# LANGUAGE TupleSections #-}
 
-module Lib (task1) where
+module Lib (task1, task2) where
 
 import Control.Monad (join)
-import Data.List (partition, (!?))
+import Data.List (partition)
 import qualified Data.Set as Set
+import qualified Data.Map.Strict as Map
+import Data.Maybe (maybeToList)
 import Flow
 
-type World = [[Char]]
 type Point = (Int, Int)
+type World = Map.Map Point Char
 type Cluster = Set.Set Point
 
 task1 :: String -> IO Int
 task1 input = let allClusters = readWorld input |> clusters
               in do
-                 return $ sum $ map price allClusters
+                 return $ sum $ map price1 allClusters
+
+task2 :: String -> IO Int
+task2 input = let allClusters = readWorld input |> clusters
+              in do
+                 return $ sum $ map price2 allClusters
 
 readWorld :: String -> World
-readWorld = lines
+readWorld input = lines input |> zip [0..] |> map (\(y, line) -> zip [0..] line |> map (\(x, char) -> ((x, y), char))) |> join |> Map.fromList
 
 clusters :: World -> [Cluster]
-clusters world = points |> addToClusters []
+clusters world = Map.keys world |> addToClusters []
   where
-    points :: [Point]
-    points = zip [0..] world |> map (\(y, line) -> take (length line) [0..] |> map (\x -> (x,y))) |> join
-
     addToClusters :: [Cluster] -> [Point] -> [Cluster]
     addToClusters existingClusters [] = existingClusters
     addToClusters existingClusters (point:restPoints) =
@@ -34,11 +38,11 @@ clusters world = points |> addToClusters []
 
     belongingToCluster :: Cluster -> Point -> Bool
     belongingToCluster cluster point =
-      let char = inWorld world point
-      in adjacentPoints point |> filter (\p -> char == inWorld world p) |> any (\x -> Set.member x cluster)
+      let char = Map.lookup point world
+      in adjacentPoints point |> filter (\p -> char == Map.lookup p world) |> any (\x -> Set.member x cluster)
 
-price :: Cluster -> Int
-price cluster = length cluster * perimiter cluster
+price1 :: Cluster -> Int
+price1 cluster = length cluster * perimiter cluster
 
 perimiter :: Cluster -> Int
 perimiter cluster = Set.toList cluster |> map perimiterOnSquare |> sum
@@ -46,13 +50,50 @@ perimiter cluster = Set.toList cluster |> map perimiterOnSquare |> sum
     perimiterOnSquare :: Point -> Int
     perimiterOnSquare position = adjacentPoints position |> filter (\pos -> Set.notMember pos cluster) |> length
 
-inWorld :: World -> Point -> Maybe Char
-inWorld world (x,y) = do
-  line <- world !? y
-  line !? x
+price2 :: Cluster -> Int
+price2 cluster = length cluster * edgeCount cluster
+
+data Direction = DUp | DRight | DDown | DLeft
+  deriving (Ord, Eq)
+data Edge = Edge Direction Point
+  deriving (Ord, Eq)
+
+offset :: Direction -> Point
+offset DUp = (0, 1)
+offset DRight = (1, 0)
+offset DDown = (0, -1)
+offset DLeft = (-1, 0)
+
+edgeCount :: Cluster -> Int
+edgeCount cluster = Set.toList cluster >>= edges |> addEdge Set.empty |> length
+  where
+    edges :: Point -> [Edge]
+    edges point = [DUp, DRight, DDown, DLeft] |> map (\direction -> edgeAt point direction |> maybeToList) |> join
+
+    edgeAt :: Point -> Direction -> Maybe Edge
+    edgeAt from@(fx, fy) direction =
+      let to@(tx, ty) = offset direction |> vAdd from
+      in if Set.member to cluster
+         then Nothing
+         else Just (Edge direction (max fx tx, max fy ty))
+
+    addEdge :: Set.Set Edge -> [Edge] -> Set.Set Edge
+    addEdge existingEdges [] = existingEdges
+    addEdge existingEdges (edge:edgesRest) =
+      if Set.member lastEdge existingEdges
+      then addEdge (Set.delete lastEdge existingEdges |> Set.insert edge) edgesRest
+      else addEdge (Set.insert edge existingEdges) edgesRest
+        where
+         lastEdge = case edge of
+           Edge DLeft (x,y) -> Edge DLeft (x,y-1)
+           Edge DRight (x,y) -> Edge DRight (x,y-1)
+           Edge DUp (x,y) -> Edge DUp (x-1,y)
+           Edge DDown (x,y) -> Edge DDown (x-1,y)
+
 
 adjacentPoints :: Point -> [Point]
-adjacentPoints point = [(-1, 0), (0, -1), (1, 0), (0, 1)] |> map (offset point)
+adjacentPoints point = [(-1, 0), (0, -1), (1, 0), (0, 1)] |> map (vAdd point)
+
+vAdd :: Num a => (a, a) -> (a, a) -> (a, a)
+vAdd (x1, y1) (x2, y2) = (x1 + x2, y1 + y2)
 
-offset :: Num a => (a, a) -> (a, a) -> (a, a)
-offset (x1, y1) (x2, y2) = (x1 + x2, y1 + y2)