Project Euler Problem 61.

pe61.hs

module Main where

import Control.Monad (guard)

import Data.Array (Ix)


main :: IO ()
main = do
  print . sum . head $ solve 3 example1 -- at least one solution exists
  print . sum . head $ solve 6 example2 -- at least one solution exists


data Figurate = Triangle | Square | Pentagonal | Hexagonal | Heptagonal | Octagonal
    deriving (Show, Read, Eq, Ord, Enum, Bounded, Ix)


example1, example2 :: [(Figurate, Int)]
example1 =
  tag Triangle   (between 1000 9999 $ triangles) ++
  tag Square     (between 1000 9999 $ squares) ++
  tag Pentagonal (between 1000 9999 $ pentagonals)
example2 =
  tag Triangle   (between 1000 9999 $ triangles) ++
  tag Square     (between 1000 9999 $ squares) ++
  tag Pentagonal (between 1000 9999 $ pentagonals) ++
  tag Hexagonal  (between 1000 9999 $ hexagonals) ++
  tag Heptagonal (between 1000 9999 $ heptagonals) ++
  tag Octagonal  (between 1000 9999 $ octagonals)


solve :: Int -> [(Figurate, Int)] -> [[Int]]
solve n xs = xs >>= step n [] [] xs


step :: Int -> [Int] -> [Figurate] -> [(Figurate, Int)] -> (Figurate, Int) -> [[Int]]
step n [] _ targets (figurate, number) = targets >>= step n [number] [figurate] targets
step n solution@(s:_) figurates targets (figurate, number) = do
  guard $ figurate `notElem` figurates && s `composes` number
  if length solution == n - 1
    then do
      guard $ number `composes` last solution
      [number:solution]
    else do
      targets >>= step n (number:solution) (figurate:figurates) targets


triangles, squares, pentagonals, hexagonals, heptagonals, octagonals :: [Int]
triangles   = for [1..] (\n -> n * (n + 1) `div` 2)
squares     = for [1..] (\n -> n * n)
pentagonals = for [1..] (\n -> n * (3 * n - 1) `div` 2)
hexagonals  = for [1..] (\n -> n * (2 * n + 1))
heptagonals = for [1..] (\n -> n * (5 * n - 3) `div` 2)
octagonals  = for [1..] (\n -> n * (3 * n - 2))


tag :: b -> [a] -> [(b, a)]
tag t = map (\x -> (t, x))

composes :: Integral a => a -> a -> Bool
a `composes` b = a `rem` 100 == b `quot` 100


for :: [a] -> (a -> b) -> [b]
for = flip map

between :: Ord a => a -> a -> [a] -> [a]
between a b = dropWhile (< a) . takeWhile (<= b)