Waterflow.hs

{-# LANGUAGE TupleSections #-}

module Waterflow
where

import Control.Applicative
import Data.Map (Map)
import qualified Data.Map as M
import Test.QuickCheck

water = sum $
  zipWith (-) 
    (zipWith min (scanl1 max h) (scanr1 max h))
    h

newtype Small a = Small {unSmall :: a}

instance Show a => Show (Small a) where
  show = show
  showsPrec = showsPrec

instance Arbitrary a => Arbitrary (Small a) where
  arbitrary = sized $ \s -> Small <$> resize (s `div` 6) arbitrary
  shrink (Small a) = Small <$> shrink a

prop_water_equal (NonEmpty ys) = water xs == trivialWater xs
  where xs :: [Int]
        xs = unSmall <$> ys

data Cell = Dirt | Water
  deriving (Eq, Show)

trivialWaterMap xs = fst $ head $ dropWhile (uncurry (/=)) (zip steps $ tail steps)
  where groundLevel = minimum xs
        ys = concat [(,i) <$> [x, x-1 .. groundLevel] | (i,x) <- zip [0..] xs]
        dirtMap = M.fromList $ map (,Dirt) ys
        (rows, cols) = (maximum xs, length xs)
        zs = [(r,c) | r <- [groundLevel .. rows], c <- [0 .. cols - 1]]
        waterMap = M.fromList $ map (,Water) zs
        initialMap = dirtMap `M.union` waterMap
        steps = iterate (flowAll zs) initialMap

trivialWater = M.foldl inc 0 . trivialWaterMap
  where inc acc v = case v of
                      Water -> acc + 1
                      Dirt -> acc

flow m loc cell = case cell of
  Dirt -> Just Dirt
  Water -> Water <$ (find (succ <$> loc) >> find (pred <$> loc))
  where find = (`M.lookup` m)

flowAll ks m = foldl (flip $ M.updateWithKey (flow m)) m ks