langs/Haskell/Lockers.hs

24 lines
760 B
Haskell

import Data.Foldable
import qualified Data.Map.Lazy as Map
type Lockers = Map.Map Int Bool
makeLockers :: Int -> Lockers
makeLockers n = Map.fromList $ zip [1..n] $ replicate n False
changeLockersState :: Int -> Lockers -> Lockers
changeLockersState n l = foldr' (Map.adjust not) l modNList
where modNList = [ n * m | m <- take (Map.size l `quot` n) [1..] ]
runLockers :: Lockers -> Lockers
runLockers l = foldr' changeLockersState l [1..Map.size l]
main :: IO ()
main = do
putStrLn "Enter how many lockers: "
n <- fmap read getLine :: IO Int
let lockers = runLockers $ makeLockers n
putStrLn $ "The following lockers remained open after " ++ show n ++ " iterations: "
print (Map.keys $ Map.filterWithKey (flip const) lockers)