ラグランジュの四平方定理 再び

id:Gemma さんのアドバイスにより、お利口にしてみた。実行時間がだいぶ縮まる。

import System
import Control.Monad

lagrange :: Integer -> [(Integer, Integer, Integer, Integer)]
lagrange x = do let seed = [ s^2 | s <- [0..(floor $ sqrt $ fromInteger x)]]
                a <- seed
                b <- [ b | b <- seed, b <= a, a+b <= x ]
                c <- [ c | c <- seed, c <= b, a+b+c <= x ]
                d <- [ d | d <- seed, d <= c, a+b+c+d == x ]
                return (a,b,c,d)
main :: IO ()
main = do (arg:_) <- getArgs
          let result = lagrange $ read arg
          mapM print result
          print $ length result

実行結果。PowerPC G4 1GH。

% time ./lagrange 20007
(5929,5929,4900,3249)
(6241,5476,5041,3249)
...
(19881,100,25,1)
(19881,121,4,1)
738
./lagrange 20007 6.05s user 0.11s system 63% cpu 9.671 total