pyth.hs

   1 -- A generator of all Pythagorean triples
   2 -- (c) 2008 Martin Mares <mj@ucw.cz>
   3
   4 -- Square
   5 sqr x = x*x
   6
   7 -- Find (p,q) such that x=p^2*q and q is square-free
   8 sqf' :: Int -> Int -> (Int,Int)
   9 sqf' _ 1 = (1,1)
  10 sqf' d x | x < (sqr d)        = (1,x)
  11          | x `mod` (sqr d) == 0 = (d*p1,q1)
  12          | x `mod` d == 0     = (p2,d*q2)
  13          | True               = sqf' (d+1) x
  14          where (p1,q1) = sqf' d (x `div` (sqr d))
  15                (p2,q2) = sqf' d (x `div` d)
  16 sqf = sqf' 2
  17
  18 -- Find all (a,b) such that c+r*a = b^2
  19 sc :: Int -> Int -> [(Int,Int)]
  20 sc c r = [ ((xx-c) `div` r, x) | x<-[1..], xx<-[sqr x], xx>c, (xx-c) `mod` r == 0 ]
  21
  22 -- Find all (a,b) such that d*(d+2a) = b^2
  23 pd d = pf (sqf d)
  24
  25 -- The same with d split to squares and square-free part
  26 pf :: (Int,Int) -> [(Int,Int)]
  27 pf (p,q) | q `mod` 2 == 0 = map (\(a,b) -> (a,2*b)) (pf' p (q `div` 2) 1)
  28          | True           = pf' p q 2
  29
  30 -- The same with an odd square-free part
  31 pf' :: Int -> Int -> Int -> [(Int,Int)]
  32 pf' p 1 r = map (\(a,b) -> (a,p*b)) (sc (sqr p) r)
  33 pf' p q r = map (\(a,b) -> (q*a,q*b)) (pf' p 1 r)
  34
  35 -- Find all Pythagorean triples with c-a = d
  36 pdt :: Int -> [(Int, Int, Int)]
  37 pdt d = map (\(a,b) -> (a,b,a+d)) (pd d)
  38
  39 -- Merging over all differences
  40 -- For every difference, we generate the first triple in time sqrt(d) and each
  41 -- next in O(1). Therefore if we add a new difference after generating d triples
  42 -- (one for each difference), the amortized cost of the initializations is O(1).
  43 pyt :: [[(Int,Int,Int)]] -> Int -> [(Int,Int,Int)]
  44 pyt l d = (map head l) ++ (pyt ((map tail l) ++ [pdt d]) (d+1))
  45 pyth' = pyt [] 1
  46
  47 -- Finally filter out duplicates
  48 -- (we always generate both (a,b,c) and (b,a,c), the one with a>b first)
  49 pyth = [ (a,b,c) | (a,b,c) <- pyth', a>b ]