A generator of Pythagorean triples.

diff --git a/pyth.hs b/pyth.hs
new file mode 100644 (file)
index 0000000..82e44fd
--- /dev/null
+++ b/pyth.hs
@@ -0,0 +1,48 @@
+-- A generator of all Pythagorean triples
+-- (c) 2008 Martin Mares <mj@ucw.cz>
+
+-- Square
+sqr x = x*x
+
+-- Find (p,q) such that x=p^2*q and q is square-free
+sqf' :: Int -> Int -> (Int,Int)
+sqf' _ 1 = (1,1)
+sqf' d x | x `mod` (sqr d) == 0 = (d*p1,q1)
+         | x `mod` d == 0     = (p2,d*q2)
+        | True               = sqf' (d+1) x
+        where (p1,q1) = sqf' d (x `div` (sqr d))
+              (p2,q2) = sqf' d (x `div` d)
+sqf = sqf' 2
+
+-- Find all (a,b) such that c+r*a = b^2
+sc :: Int -> Int -> [(Int,Int)]
+sc c r = [ ((xx-c) `div` r, x) | x<-[1..], xx<-[sqr x], xx>c, (xx-c) `mod` r == 0 ]
+
+-- Find all (a,b) such that d*(d+2a) = b^2
+pd d = pf (sqf d)
+
+-- The same with d split to squares and square-free part
+pf :: (Int,Int) -> [(Int,Int)]
+pf (p,q) | q `mod` 2 == 0 = map (\(a,b) -> (a,2*b)) (pf' p (q `div` 2) 1)
+         | True           = pf' p q 2
+
+-- The same with an odd square-free part
+pf' :: Int -> Int -> Int -> [(Int,Int)]
+pf' p 1 r = map (\(a,b) -> (a,p*b)) (sc (sqr p) r)
+pf' p q r = map (\(a,b) -> (q*a,q*b)) (pf' p 1 r)
+
+-- Find all Pythagorean triples with c-a = d
+pdt :: Int -> [(Int, Int, Int)]
+pdt d = map (\(a,b) -> (a,b,a+d)) (pd d)
+
+-- Merging over all differences
+-- For every difference, we generate the first triple in time sqrt(d) and each
+-- next in O(1). Therefore if we add a new difference after generating d triples
+-- (one for each difference), the amortized cost of the initializations is O(1).
+pyt :: [[(Int,Int,Int)]] -> Int -> [(Int,Int,Int)]
+pyt l d = (map head l) ++ (pyt ((map tail l) ++ [pdt d]) (d+1))
+pyth' = pyt [] 1
+
+-- Finally filter out duplicates
+-- (we always generate both (a,b,c) and (b,a,c), the one with a>b first)
+pyth = [ (a,b,c) | (a,b,c) <- pyth', a>b ]