Project Euler 98

http://projecteuler.net/index.php?section=problems&id=98


ペアを作って、桁数が合う平方数すべてについてマッチするか調べる。

import qualified Data.Map as M
import Data.List

int_sqrt n = floor (sqrt (fromIntegral n))
is_square n = m * m == n where m = int_sqrt n

combinations a 0 = [[]]
combinations a n = concat (map (\e -> map (e:)
            (combinations (filter (>e) a) (n - 1))) a)

digits 0 = []
digits n = (digits (div n 10)) ++ [mod n 10]
to_number a = foldl (\x y -> x * 10 + y) 0 a

get_words [] s = []
get_words (',':cs) s = get_words cs []
get_words ('"':cs) [] = get_words cs []
get_words ('"':cs) s = s:(get_words cs [])
get_words (c:cs) s = get_words cs (s ++ [c])

ins k v m | M.member k m = M.insert k ((m M.! k) ++ [v]) m
          | otherwise   = M.insert k [v] m
collect :: [[Char]] -> M.Map [Char] [[Char]] -> M.Map [Char] [[Char]]
collect [] m = m
collect (s:ss) m = collect ss (ins (sort s) s m)

perm_pairs m = concat [ combinations a 2 | a <- map snd (M.toList m) ]
match [] [] a = a
match (c:cs) (d:ds) a = if p == -1 then []
                    else match cs ds (if p == 1 then a else (c,d):a) where
        p = f c d a
        f c d [] = 0
        f c d ((p,q):xs) = if p == c && q == d then 1
                      else if p == c || q == d then -1
                                               else f c d xs

match2 [] a = []
match2 (c:cs) a = if d == -1 then [] else d:(match2 cs a) where
        d = f c a
        f c [] = -1
        f c ((p,q):xs) | c == p    = q
                       | otherwise = f c xs

match3 [s1,s2] = f s1 s2 (map (^2) [s..e-1]) where
        l = length s1
        (s,e) = if even l then ((int_sqrt (10^(l-1))) + 1, 10^(div l 2))
                          else (10^(div l 2), (int_sqrt (10^l)) + 1)
        f s1 s2 [] = 0
        f s1 s2 (n:ns) | null a      = f s1 s2 ns
                       | m == 0      = f s1 s2 ns
                       | is_square m = max m n
                       | otherwise   = f s1 s2 ns where
                a = match s1 (digits n) []
                b = match2 s2 a
                m = if head b == 0 then 0 else to_number b

solve cs = map match3 (perm_pairs (collect (get_words cs []) M.empty))

main = do
    cs <- readFile "words.txt"
    print (foldr max 0 (solve cs))