set-cover-0.1.2: example/PythagorasColoring.hs

{- |
Assign one of two colors to every natural number
such that there is no uniformly colored Pythagorean triple.

<https://en.wikipedia.org/wiki/Boolean_Pythagorean_triples_problem>
<https://dorfuchs93.github.io/boolean-pythagorean-triples/>
-}
module Main where

import qualified Math.SetCover.Exact as ESC

import qualified Data.Traversable as Trav
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import qualified Data.Set as Set
import Data.IntMap (IntMap)
import Data.IntSet (IntSet)
import Data.Set (Set)
import Data.Maybe (catMaybes)

import Control.Monad (guard)
import Control.Applicative (liftA2)


data Color = Red | Blue
   deriving (Eq, Ord, Show)

data X =
        Number Int
      | PythagoreanTriple IntSet | PythagoreanColor IntSet Int Color
   deriving (Eq, Ord, Show)

type Assign = ESC.Assign (Maybe (Int, Color)) (Set X)


colorNumber :: [IntSet] -> Int -> Color -> Assign
colorNumber tuples n c =
   ESC.assign (Just (n,c)) $ Set.fromList $
   Number n
   :
   map (\tuple -> PythagoreanColor tuple n c)
      (filter (IntSet.member n) tuples)

constIntMap :: a -> IntSet -> IntMap a
constIntMap x = IntMap.fromSet (const x)

admissibleTupleColorings :: IntSet -> [IntMap Color]
admissibleTupleColorings xs =
   filter
      (\tuple ->
         tuple /= constIntMap Red xs
         &&
         tuple /= constIntMap Blue xs) $
   Trav.sequence $ constIntMap [Red,Blue] xs


assigns :: [IntSet] -> [Assign]
assigns tuples =
   let omega = IntSet.unions tuples in
   liftA2 (colorNumber tuples) (IntSet.toList omega) [Red, Blue]
   ++
   (do
      tuple <- tuples
      tupleColoring <- admissibleTupleColorings tuple
      [ESC.assign Nothing $ Set.fromList $
         PythagoreanTriple tuple :
         IntMap.elems
            (IntMap.mapWithKey (PythagoreanColor tuple) tupleColoring)])


intMapFromLabels :: [Maybe (Int, Color)] -> IntMap Color
intMapFromLabels = IntMap.fromList . catMaybes

pythagoreanTriples :: [IntSet]
pythagoreanTriples =
   IntSet.fromList [3,4,5] :
   IntSet.fromList [5,12,13] :
   IntSet.fromList [9,12,15] :
   IntSet.fromList [12,16,20] :
   []


sqr :: Num a => a -> a
sqr a = a*a

{-
Alternatively we could use Euclidean's formula.
-}
pythagoreanTriplesUpTo :: Int -> [IntSet]
pythagoreanTriplesUpTo n = do
   a <- [1..n]
   b <- [a+1..n]
   let a2b2 = sqr a + sqr b
   let c = round $ sqrt (fromIntegral a2b2 :: Double)
   guard $ c<=n && a2b2 == sqr c
   return $ IntSet.fromList [a,b,c]


mainAll :: IO ()
mainAll =
   mapM_ (print . intMapFromLabels) $
   ESC.partitions $ assigns pythagoreanTriples

main :: IO ()
main =
   mapM_ (print . intMapFromLabels) $ take 1 $
   ESC.partitions $ ESC.intSetFromSetAssigns $
   assigns $ pythagoreanTriplesUpTo 1500