Sudoku.hs

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wall -Werror #-}

module Sudoku where

import Data.Char (intToDigit)
import Data.List (intercalate, sortOn)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.List.Split (chunksOf)
import Data.Maybe (maybeToList)

type Value = Int

type Cell = NonEmpty Value

type Pos = Int

type Puzzle = [(Pos, Cell)]

cells :: Puzzle -> [Cell]
cells = map snd . sortOn fst

getXY :: Pos -> (Int, Int)
getXY n = (n `mod` 9, n `div` 9)

conflictsWith :: Pos -> Pos -> Bool
conflictsWith p p' =
  let (x, y) = getXY p
      (x', y') = getXY p'
   in x == x' || y == y' || (x `div` 3 == x' `div` 3 && y `div` 3 == y' `div` 3)

constrain :: Pos -> Value -> Puzzle -> Maybe Puzzle
constrain focus v =
  mapM (\(p, c) -> (p,) <$> if p `conflictsWith` focus then restrict c else pure c)
  where
    restrict = NE.nonEmpty . NE.filter (/= v)

sortByComplexity :: Puzzle -> Puzzle
sortByComplexity = sortOn (NE.length . snd)

solve :: Puzzle -> Puzzle
solve = head . search
  where
    search (sortByComplexity -> (p, vs) : cs) = do
      v <- NE.toList vs
      cs' <- maybeToList (constrain p v cs)
      ((p, v :| []) :) <$> search cs'
    search _ = [[]]

-- IO

pretty :: Puzzle -> String
pretty = prettyLines . map prettyCell . cells
  where
    prettyCell (x :| []) = intToDigit x
    prettyCell _ = ' '
    prettyLines =
      ("-------------------\n" ++)
        . (++ "\n-------------------")
        . intercalate "|-----+-----+-----|\n"
        . chunksOf 60
        . intercalate "\n"
        . map (("| " ++) . (++ " |") . intercalate " | " . chunksOf 3)
        . chunksOf 9

readPuzzle :: String -> Puzzle
readPuzzle = zip [0 ..] . map (makeCell . read . pure)
  where
    makeCell 0 = 1 :| [2 .. 9]
    makeCell n = n :| []

-- Testing

testPuzzleEasy :: Puzzle
testPuzzleEasy =
  readPuzzle
    "040000179\
    \002008054\
    \006005008\
    \080070910\
    \050090030\
    \019060040\
    \300400700\
    \570100200\
    \928000060"

testPuzzleHard :: Puzzle
testPuzzleHard =
  readPuzzle
    "000003017\
    \015009008\
    \060000000\
    \100007000\
    \009000200\
    \000500004\
    \000000020\
    \500600340\
    \340200000"

testPuzzleHardest :: Puzzle
testPuzzleHardest =
  readPuzzle
    "100007090\
    \030020008\
    \009600500\
    \005300900\
    \010080002\
    \600004000\
    \300000010\
    \040000007\
    \007000300"

solves :: Puzzle -> Puzzle -> Bool
solves p1 p2 = and $ zipWith solvesCell (cells p1) (cells p2)
  where
    solvesCell (x :| []) xs = x `elem` xs
    solvesCell _ _ = False

runTests :: Bool
runTests =
  solve testPuzzleEasy `solves` testPuzzleEasy
    && solve testPuzzleHard `solves` testPuzzleHard
    && solve testPuzzleHardest `solves` testPuzzleHardest