Commit cb1f150

@@ -6,9 +6,16 @@ describe "problem nine" do
 
   # There exists exactly one Pythagorean triplet for which a + b + c = 1000.
   # Find the product abc.
-  #
+
+  # a < b < c
+  # c^2 = a^2 + b^2
   class PythagoreanTriplet
     include Enumerable
+    attr_reader :max
+
+    def initialize(max = 1_000)
+      @max = max
+    end
 
     def each(&block)
       triplets.each(&block)
@@ -18,11 +25,12 @@ describe "problem nine" do
 
     def triplets
       Enumerator.new do |yielder|
-        n = 0
-        loop do
-          x, y, z = n*n, (n+1)*(n+1), (n+2)*(n+2)
-          yielder.yield([n, n+1, n+2]) if x + y == z
-          n = n+1
+        (2..max).each do |m|
+          (1...m).each do |n|
+            # https://en.wikipedia.org/wiki/Pythagorean_triple#Proof_of_Euclid.27s_formula
+            a, b, c = m*m - n*n, 2*m*n, m*m + n*n
+            yielder.yield([a, b, c])
+          end
         end
       end
     end
@@ -37,12 +45,14 @@ describe "problem nine" do
       (triplet.last * triplet.last) == 25
     end
     expect(result).to eql([3, 4, 5])
+    expect(subject).to include([3, 4, 5])
+    expect(subject).to include([5, 12, 13])
   end
 
   it "returns the triplet for 1000" do
     result = subject.find do |triplet|
-      (triplet.last * triplet.last) == 1_000
+      (triplet[0] + triplet[1] + triplet[2]) == 1_000
     end
-    expect(result).to eql(10)
+    expect(result).to eql([375, 200, 425])
   end
 end