Day 9, still not right

advent-of-code.cabal

@@ -8,7 +8,7 @@ author:             Stuart A. Kurtz
 maintainer:         stuart@cs.uchicago.edu
 
 common aoc-common
-  build-depends       : base ^>=4.14.3.0
+  build-depends       : base ^>=4.14.3.0, array, containers
   hs-source-dirs      : src
   other-extensions    : TypeApplications
   default-language    : Haskell2010
@@ -98,7 +98,6 @@ executable aoc-09b
 executable aoc-10a
   import              : aoc-common
   main-is             : AOC-10a.hs
-  build-depends       : array
 
 executable aoc-10b
   import              : aoc-common

src/AOC-09b.hs

@@ -2,6 +2,8 @@ module Main where
 
 import Data.Array
 import Data.List
+import qualified Data.Set as S
+import Data.Set (Set)
 
 makeHeightMap :: String -> Array (Int,Int) Int
 makeHeightMap s = array ((1,1),(mx,my)) entries where
@@ -10,7 +12,7 @@ makeHeightMap s = array ((1,1),(mx,my)) entries where
                        , (y,c) <- zip [1..] row
                        , let z = read [c]
                        ]
-  mx = length (rows)
+  mx = length rows
   my = length (head rows)
 
 heightAt :: Array (Int,Int) Int -> (Int,Int) -> Int
@@ -34,19 +36,32 @@ lowPoints arr = filter (isLowPoint arr) . range . bounds $ arr
 -- This code is wrong for the interpretation of a basin given by
 -- the problem.
 
-basin :: Array (Int,Int) Int -> (Int,Int) -> [(Int,Int)]
-basin arr pos = pos : iter pos (neighbors arr pos) where
-  iter :: (Int,Int) -> [(Int,Int)] -> [(Int,Int)]
-  iter _ [] = []
-  iter p (n:ns) 
-    | n `flowsTo` p = n : iter p ns ++ iter n (neighbors arr n)
-    | otherwise = iter p ns
-  flowsTo n p =
-    let target = heightAt arr n
-        flowPoints = filter ((< target) . (arr !)) . neighbors arr $ n
-    in flowPoints == [p]
+basin :: Array (Int,Int) Int -> (Int,Int) -> Set (Int,Int)
+basin arr pos = search (S.singleton pos) where
+  frontier :: Set (Int,Int) -> Set (Int,Int)
+  frontier pts = expansion `S.difference` pts where
+    expansion = S.fromList . foldr (\pt rest -> neighbors arr pt ++ rest) [] $ pts
+  search :: S.Set (Int,Int) -> S.Set (Int,Int)
+  search pts
+    | not . S.null $ newpoints = search (pts `S.union` newpoints)
+    | otherwise = pts
+    where
+      newpoints = S.fromList 
+                  [pt | pt <- S.toList (frontier pts)
+                      , flowPoint pt
+                      , S.fromList (flowsFrom pt) `S.isSubsetOf` pts
+                  ]
+      flowsFrom pt = [p | p <- neighbors arr pt
+                        , arr ! p < arr ! pt]
+      flowPoint pt = arr ! pt /= 9 && all ((/= arr ! pt) . (arr !)) (neighbors arr pt)
 
 main :: IO ()
 main = do
-  heightMap <- makeHeightMap <$> readFile "test/heightMap.txt"
-  putStr . unlines . map (show . sort) $ basin heightMap <$> lowPoints heightMap
+  heightMap <- makeHeightMap <$> readFile "data/heightMap.txt"
+  putStr . show . take 3 . reverse . sort . map S.size $ basin heightMap <$> lowPoints heightMap
+
+{- failed answers:
+
+414960
+
+-}
\ No newline at end of file