Percolation solution in Haskell

README.rst

To test this in ghci:

import Percolation
> percolated $ addSites [(2,2),(3,3),(0,1),(0,2),(0,3),(0,4)] $ newGrid 5
False
> percolated $ addSites [(2,2),(3,3),(0,1),(0,2),(0,3),(0,4),(0,0)] $ newGrid 5
True

Percolation.hs

module Percolation where
                                                                                
import Control.Monad (join)
import Data.Maybe (mapMaybe)
import qualified Data.Map.Lazy as M
import qualified Data.Set as S

-- | Percolation algorithm.
--
-- Important assumptions or implementation details:
--
-- * When sets are unioned, the lowest SetId is used as
--   the SetId for the new set.  This allows us to check
--   whether we have percolated by asking whether Set 1
--   (initially the bottom) has been removed due to merging
--   with Set 0 (initially the top).

type SiteId = Int
type SetId = Int
type Coord = (Int,Int)

data Grid =
  Grid
    Int 
    -- ^ Dimension of matrix
    SiteId
    -- ^ Next SiteId to allocate
    (M.Map (Int,Int) (Maybe SiteId))
    -- ^ Mapping of coordinate to SiteId
    (M.Map SiteId SetId)
    -- ^ Mapping of SiteId to SetId
    (S.Set SetId)
    -- ^ Mapping of SetId to set of connected coordinates
  deriving (Show)
  
newGrid :: Int -> Grid
newGrid n =
  let
    coords = [ (x,y) | x <- [0..n-1], y <- [0..n-1] ]
    topCoords = [ (x,-1) | x <- [0..n-1] ]
    bottomCoords = [ (x,n) | x <- [0..n-1] ]
    coordMap = foldr (`M.insert` Nothing) M.empty coords
    coordMap' = foldr (`M.insert` Just 0) coordMap topCoords
    coordMap'' = foldr (`M.insert` Just 1) coordMap' bottomCoords
  in
    Grid n 2 coordMap''
      (M.fromList [(0,0),(1,1)])
      (S.fromList [0, 1])

percolated :: Grid -> Bool
percolated (Grid _ _ _ _ sets) = S.notMember 1 sets

addSite :: Coord -> Grid -> Grid
addSite xy@(x,y) g@(Grid n nextSiteId siteIds setIds sets) =
  case M.lookup xy siteIds of
    Just Nothing ->
      -- site not open; find surrounding setIds
      let
        adjSiteIds = mapMaybe (join . (`M.lookup` siteIds))
          [(x,y-1),(x-1,y),(x+1,y),(x,y+1)]
        adjSetIds = mapMaybe (`M.lookup` setIds) adjSiteIds
        siteId = foldr min nextSiteId adjSiteIds
        setId = foldr min nextSiteId adjSetIds
      in
        Grid n
          (if setId == nextSiteId then nextSiteId + 1 else nextSiteId)
          (M.insert xy (Just siteId) siteIds)
          (foldr (`M.insert` setId) setIds (setId : adjSetIds))
          (S.insert setId
            $ foldr S.delete sets adjSetIds)
    _ ->
      -- site out of bounds, or already open; return grid unchanged
      g

addSites :: [Coord] -> Grid -> Grid
addSites = foldr ((.) . addSite) id