HaskellでGoogle Code Jamの数独チェッカー問題を解いてみました。

問題

Google Code Jam の
Google of Greater China Test for New Grads of 2014のRound B China New Grad Test 2014のProblem A. Sudoku Checkerを解いてみました。

問題の概要

問題の概要は、数独の解答が与えられて、正しいかどうか判断する問題。

問題の入出力

問題の入出力は
入力 -> 数独の解答
出力 -> 解答が正しいときYes, 間違っているときNo

実装

そして、これが解いたコード。

import Control.Monad
import Text.Printf
import Control.Applicative
import Data.List

main :: IO()
main = do
  testCase <- readLn
  forM_ [1..testCase] $ \i -> do
    n <- read <$> getLine
    mt <- replicateM (n*n :: Int) $ map read . words <$> getLine
    printf "Case #%d: %s\n" (i :: Int) $ slv n mt

slv :: Int -> [[Int]] -> String
slv n mt | rowOK && colOK && subMtOK = "Yes"
         | otherwise = "No"
  where rowOK = all contNum mt
        colOK = all contNum $ transpose mt
        subMtOK = all contNum subMt
        subMt = marN n (concatMap transpose (divN n mt [])) []

contNum :: [Int] -> Bool
contNum ls = and $ zipWith (==) (sort ls) [1..]

divN :: Int -> [[Int]] -> [[[Int]]] -> [[[Int]]]
divN _ [] ret = ret
divN n ls ret = divN n l' (l:ret)
  where (l,l') = splitAt n ls

marN :: Int -> [[Int]] -> [[Int]] -> [[Int]]
marN _ [] ret = ret
marN n ls ret = marN n l' (concat l:ret)
  where (l,l') = splitAt n ls

コードの内容

main関数

main関数は問題内容を読み込んで、slv関数で問題を解かせて答えを出力。

slv関数

数独の解答が正しいかどうかの条件は

• 各行に関して、全ての数字が存在する(例えば9×9の大きい正方形に対して1～9まで存在する)
• 各列に関して、全ての数字が存在する
• 小さい正方形に対して全ての数字が存在する

である。

したがって、上記の条件を全て満たしているとき"Yes"
一つでも満たさなければ"No"
としてあげればよい。

rowOK関数

数独の解答を渡して、各行に全ての数字が存在しているかチェックしている
存在しているとTrue, 存在してないとFalse

colOK関数

数独の解答を渡して、各列に全ての数字が存在しているかチェックしている
存在しているとTrue, 存在してないとFalse

subMtOK関数

数独の解答を渡して、各小さい正方形に全ての数字が存在しているかチェックしている
存在しているとTrue, 存在してないとFalse

contNum関数

これは全ての数字が存在しているかをチェックしている。
問題内容のサイズが小さいから、どうでも良いかもしれないけど、
普通にintersectionするよりも
sortしてから数字を照らし合わせたほうが計算量のオーダーは小さいはず。
大きい正方形のサイズを$m$とすると

• intersect関数は$O(m^2)$
• contNum関数は$O(mlogm)$

($n$だと問題文と混ざるので$m$で表しました。)

divN関数 marN関数 subMt関数(小さい正方形の取得方法)

今回は以下のように変換した。
以下は例題の問題1の例

[[5,3,4,6,7,8,9,1,2],
 [6,7,2,1,9,5,3,4,8],
 [1,9,8,3,4,2,5,6,7],
 [8,5,9,7,6,1,4,2,3],
 [4,2,6,8,5,3,7,9,1],
 [7,1,3,9,2,4,8,5,6],
 [9,6,1,5,3,7,2,8,4],
 [2,8,7,4,1,9,6,3,5],
 [3,4,5,2,8,6,1,7,9]]
-- 問題内容を
[[[5,3,4,6,7,8,9,1,2],
  [6,7,2,1,9,5,3,4,8],
  [1,9,8,3,4,2,5,6,7]],
 [[8,5,9,7,6,1,4,2,3],
  [4,2,6,8,5,3,7,9,1],
  [7,1,3,9,2,4,8,5,6]],
 [[9,6,1,5,3,7,2,8,4],
  [2,8,7,4,1,9,6,3,5],
  [3,4,5,2,8,6,1,7,9]]]
-- と3つの行に区切ったあと(div関数の処理)
-- 縦に短冊に区切る。(concatMap transpose)
[[5,6,1],
 [3,7,9],
 [4,2,8],
 [6,1,3],
 [7,9,4],
 [8,5,2],
 [9,3,5],
 [1,4,6],
 [2,8,7],
 [8,4,7],
 [5,2,1],
 [9,6,3],
 [7,8,9],
 [6,5,2],
 [1,3,4],
 [4,7,8],
 [2,9,5],
 [3,1,6],
 [9,2,3],
 [6,8,4],
 [1,7,5],
 [5,4,2],
 [3,1,8],
 [7,9,6],
 [2,6,1],
 [8,3,7],
 [4,5,9]]
-- このリストを3つずつ要素を取得すれば小さい正方形の要素のリストが得られる。(marN関数)
[[5,6,1,3,7,9,4,2,8],
 [6,1,3,7,9,4,8,5,2],
 [9,3,5,1,4,6,2,8,7],
 [8,4,7,5,2,1,9,6,3],
 [7,8,9,6,5,2,1,3,4],
 [4,7,8,2,9,5,3,1,6],
 [9,2,3,6,8,4,1,7,5],
 [5,4,2,3,1,8,7,9,6],
 [2,6,1,8,3,7,4,5,9]]

もっと簡単に小さい正方形の要素のリストを取得したいんだけど。
これを実装した関数が、divN関数、marN関数、subMt関数である。
この小さい正方形の要素のリストを取得するのに約10行くらい使っているので
もう少しコンパクトに書けたら・・・と思った。

疑問

小さい正方形に関して、数字が全て存在しているかチェックする方法で
簡単に処理する方法がわからなかった。
(小さい正方形のリストを簡単に取り出す方法が分からなかった)

誰か簡単に取得する方法を教えてください・・・。
どうも無理やりな感じがして、この方法が最適とは思えないんだよなぁ・・・。

小さい正方形のリストを取り出さず、
Arrayを使って添え字でアクセスしてチェックすればいいのかもしれないけど、
今回はその方法を取らなかった。
(このアルゴリズムだと手続き型言語のほうが実装しやすそうなため)