Finding the lowest perfect square that can be added to a number to make another perfect square

Question

The code below returns the smallest perfect square that can be added to a number n to result in a perfect square. It works perfectly fine but I need it to run faster.

def solve n
  (1..n).each do |i|
    i = i**2 + n
    return i -= n if (Math.sqrt(i) % 1).zero?
  end
  -1
end

p solve(13) #36
 # because 36 is the smallest perfect square that can be added to 13 to form a perfect square => 13 + 36 = 49

p solve(3) #= 1 # 3 + 1 = 4, a perfect square
p solve(12) #= 4 # 12 + 4 = 16, a perfect square
p solve(9) #= 16 
p solve(4) #= -1`
 

Answer 1

I can see a few changes you can do in your code.
First one, save the calculated square and use it in the next iteration:

def solve2(n)
  # Start with an initial zero value
  pow = 0
  (1..n).each do |i|
    # This is the same as i**2, but using the square previously calculated.
    # Google squares and pascal triangle to more insight about this, but
    # in terms of cpu usage, this is faster.
    pow = pow + (i - 1) * 2 + 1
    r = pow + n
    return r -= n if (Math.sqrt(r) % 1).zero?
  end
  -1
end

And yes, it's not a big improvement, but it works. Trying in a worst case scenario, where no value is found so it must check every different iteration:

irb(main):250:0> require 'benchmark'
=> false
irb(main):251:0>
irb(main):252:0> Benchmark.bm do |x|
irb(main):253:1*   x.report { solve 50_000_002 }
irb(main):254:1>   x.report { solve2 50_000_002 }
irb(main):255:1> end
       user     system      total        real
   9.072362   0.029596   9.101958 (  9.783397)
   8.491063   0.030566   8.521629 (  9.159671) # ~600ms faster than the previous one, yay!

Next thing, instead of iterating between 1 and n, you need to iterate only integer values between Math.sqrt(1**2 + n) and Math.sqrt(n**2 + n). For example, using the current code with the same 50_000_002 value of the benchmark:

irb(main):358:0> Math.sqrt(1**2 + 50_000_002)
=> 7071.068023997506 # 1st iteration, for sure this will fail with ( % 1).zero?
irb(main):359:0> Math.sqrt(2**2 + 50_000_002)
=> 7071.068236129531 # 2nd one still with decimals
irb(main):360:0> Math.sqrt(3**2 + 50_000_002)
=> 7071.068589682892 # and so on...
irb(main):361:0> Math.sqrt(4**2 + 50_000_002)
=> 7071.069084657567
irb(main):362:0> Math.sqrt(5**2 + 50_000_002)
=> 7071.069721053526

as you can see, for the first values of the iteration, you're getting pretty much the same integer value with only decimals of difference. Eventually you'll be close to get an integer value much further in the 114th iteration:

irb(main):392:0> Math.sqrt(114**2 + 50_000_002)
=> 7071.986849535285
irb(main):393:0> Math.sqrt(115**2 + 50_000_002)
=> 7072.003040157718 # that was close, but not enough to get a integer. Keep trying...

So checking all those 114 values were useless, because their results weren't even integers. After a while, reaching the last values, you still don't get a single integer value:

irb(main):394:0> Math.sqrt(50_000_000**2 + 50_000_002)
=> 50000000.500000015
irb(main):395:0> Math.sqrt(50_000_001**2 + 50_000_002)
=> 50000001.50000001
irb(main):396:0> Math.sqrt(50_000_002**2 + 50_000_002)
=> 50000002.5

Then, the problem was to check float values when we need to check only when values have no decimals. In other words, instead of iterate between 1 and 50_000_002, you must iterate between 7072 (ceiling value for result of Math.sqrt(1**2 + 50_000_002), taken from your first previous iteration) and 50_000_002 (floor result of Math.sqrt(50_000_002**2 + 50_000_002), your last iteration, which at the end is the same value as n).
Why using this new perspective? mostly because Math.sqrt is an expensive operation compared with +, - or * (or even ** if the solve2 improvement hasn't been applied). I'll try to explain as much as I can in the code:

def solve3(n)
  # Lowest square root result. This will be the starting point 
  lowest_sqrt = Math.sqrt(1 + n).ceil
  # then getting first i value to compare. This is the last time using
  # Math.sqrt
  i = Math.sqrt(lowest_sqrt**2 - n).to_i
  # Keeping a flag, which will be the value to return later
  flag = i
  # Initial values for perfect square, where will be used the base value from
  # previous iteration (that's why i - 1) to make the power replacement as in
  # solve2
  pow = (i - 1)**2
  # and result after adding the n value.
  res = sq(pow, i) + n
  loop do
    pow = sq(pow, i)
    # When square and result are the same, that's our lowest perfect square
    return flag**2 if res == pow

    # In case result is lower, result must be recalculated with new square
    # value.
    if res < pow
      flag = i
      res = pow + n
    end
    
    i += 1
    break if i > n
  end
  -1
end

# pascal triangle trick
def sq(pow, i)
  pow + (i - 1) * 2 + 1
end

This is a significant improvement compared with previous implementations:

irb(main):908:0> Benchmark.bm do |x|
irb(main):909:1*   x.report { solve 50_000_002 }
irb(main):910:1>   x.report { solve2 50_000_002 }
irb(main):911:1>   x.report { solve3 50_000_002 }
irb(main):912:1> end
       user     system      total        real
   9.077371   0.036185   9.113556 ( 10.015861)
   8.486830   0.026608   8.513438 (  9.090664)
   4.709684   0.012402   4.722086 (  4.925373)

And I think there's still space to improve the code. I have the feeling there are some useless comparisons in the loop that can be avoided, but I've spent a while explaining all this, so I'll leave that to you :)