Day twelve - actually performant clustering

day-12-haskell/src/Lib.hs

@@ -1,9 +1,8 @@
 {-# LANGUAGE TupleSections #-}
 
-module Lib (task1, task2) where
+module Lib where
 
 import Control.Monad (join)
-import Data.List (partition)
 import qualified Data.Set as Set
 import qualified Data.Map.Strict as Map
 import Data.Maybe (maybeToList)
@@ -26,20 +25,47 @@ task2 input = let allClusters = readWorld input |> clusters
 readWorld :: String -> World
 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 = Map.keys world |> addToClusters []
-  where
-    addToClusters :: [Cluster] -> [Point] -> [Cluster]
-    addToClusters existingClusters [] = existingClusters
-    addToClusters existingClusters (point:restPoints) =
-      case partition (\c -> belongingToCluster c point) existingClusters of
-        ([], other) -> addToClusters (Set.singleton point : other) restPoints
-        (matching, other) -> addToClusters ((map Set.toList matching |> join |> Set.fromList |> Set.insert point) : other) restPoints
+data ClusterCollector = ClusterCollector !Int !(Map.Map Point Int) !(Map.Map Int Cluster)
+  deriving (Show)
+ccCreate :: ClusterCollector
+ccCreate = ClusterCollector 0 Map.empty Map.empty
+ccGetId :: Point -> ClusterCollector -> Maybe Int
+ccGetId point (ClusterCollector _ idsByPoint _) = idsByPoint Map.!? point
+ccAdd :: Cluster -> ClusterCollector -> ClusterCollector
+ccAdd cluster (ClusterCollector counter idsByPoint clustersById) = ClusterCollector (counter+1) (foldr (\p -> Map.insert p counter) idsByPoint cluster) (Map.insert counter cluster clustersById)
+ccAddTo :: Int -> Point -> ClusterCollector -> ClusterCollector
+ccAddTo into value cc@(ClusterCollector counter idsByPoint clustersById) =
+  let newCluster = clustersById Map.!? into |> expect ("Into " ++ (show into) ++ " not found in ccAddTo: " ++ show cc) |> Set.insert value
+  in ClusterCollector counter (Map.insert value into idsByPoint) (Map.insert into newCluster clustersById)
+ccCombine :: Int -> Int -> ClusterCollector -> ClusterCollector
+ccCombine into from (ClusterCollector counter idsByPoint clustersById)
+  | into == from = error "Into must be different from from in ccCombine"
+  | otherwise =
+  let fromCluster = expect "From not found in ccCombine" $ clustersById Map.!? from
+      toCluster = expect "Into not found in ccCombine" $ clustersById Map.!? into
+  in (ClusterCollector
+      counter
+      (foldr (\p -> Map.insert p into) idsByPoint fromCluster)
+      (Map.insert into (Set.union toCluster fromCluster) clustersById |> Map.delete from)
+    )
+ccClusters :: ClusterCollector -> [Cluster]
+ccClusters (ClusterCollector _ _ clustersById) = Map.elems clustersById
 
-    belongingToCluster :: Cluster -> Point -> Bool
-    belongingToCluster cluster point =
-      let char = Map.lookup point world
-      in adjacentPoints point |> filter (\p -> char == Map.lookup p world) |> any (\x -> Set.member x cluster)
+clusters :: World -> [Cluster]
+clusters world = Map.keys world |> addToClusters ccCreate |> ccClusters
+  where
+    addToClusters :: ClusterCollector -> [Point] -> ClusterCollector
+    addToClusters existingClusters [] = existingClusters
+    addToClusters existingClusters (point:restPoints) = addToClusters newClusters restPoints
+      where
+        newClusters = case adjacentEqualPoints |> map (\p -> ccGetId p existingClusters |> maybeToList) |> join |> distinct of
+          [] -> ccAdd (Set.singleton point) existingClusters
+          [match] -> ccAddTo match point existingClusters
+          (first:other) -> foldr (\o -> ccCombine first o) existingClusters other |> ccAddTo first point
+
+        adjacentEqualPoints =
+          let char = world Map.! point
+          in adjacentPoints point |> filter (\p -> world Map.!? p |> maybe False ((==) char))
 
 price1 :: Cluster -> Int
 price1 cluster = length cluster * perimiter cluster
@@ -97,3 +123,10 @@ 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)
 
+expect :: String -> Maybe a -> a
+expect _ (Just a) = a
+expect e Nothing = error e
+
+distinct :: Ord a => [a] -> [a]
+distinct = Set.toList . Set.fromList
+