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