More efficient way to bruteforce finding solutions to (x+y)^2=str(x)+str(y)? Can it be vectorised?

Question

So far I have written:

n=1000
solutions=[]
for i in range(1,n+1):
    for j in range(1,n+1):
        if str((i+j)**2)==str(i)+str(j):
            solutions.append("("+str(i)+"+"+str(j)+")^2 = "+str((i+j)**2))
for solution in solutions:
    print(solution)

This takes 1.03 seconds on my computer. Is there a quicker way to implement the comparison? I found a page on vectorisation but I'm not sure how I would generate the list I would need to then vectorise the comparison.

Answer 1

This can be accomplished even faster by searching for an (x, y) pair that satisfies the equation for a given square in your target range. In fact, this reduces the problem from O(n^2) to O(nlogn) time complexity.

def split_root(n):
    div = 10
    while div < n:
        x, y = divmod(n, div)
        div *= 10
        if not y or y < div // 100: continue
        if (x + y) ** 2 == n: yield x, y

Then just iterate over the possible squares:

def squares(n):
    for i in range(n):
        for sr in split_root(i ** 2):
            yield "({}+{})^2 = {}".format(*sr, sum(sr)**2)

Example usage:

print("\n".join(squares(100000)))

Output:

(8+1)^2 = 81
(20+25)^2 = 2025
(30+25)^2 = 3025
(88+209)^2 = 88209
(494+209)^2 = 494209
(494+1729)^2 = 4941729
(744+1984)^2 = 7441984
(2450+2500)^2 = 24502500
(2550+2500)^2 = 25502500
(5288+1984)^2 = 52881984
(6048+1729)^2 = 60481729
(3008+14336)^2 = 300814336
(4938+17284)^2 = 493817284
(60494+17284)^2 = 6049417284
(68320+14336)^2 = 6832014336

For comparison, your original solution-

def op_solver(n):
    solutions = []
    for i in range(1,n+1):
        for j in range(1,n+1):
            if str((i+j)**2)==str(i)+str(j):
                solutions.append("("+str(i)+"+"+str(j)+")^2 = "+str((i+j)**2))
    return solutions


>>> timeit("op_solver(1000)", setup="from __main__ import op_solver", number=5) / 5
0.8715057126013562

My solution-

>>> timeit("list(squares(2000))", setup="from __main__ import squares", number=100) / 100
0.006898956680088304

Roughly a 125x speedup for your example usage range, and it will run asymptotically faster as n grows.

This also has the benefit of being faster and simpler than the numpy solution, without of course requiring numpy. If you do need a faster version, I'm sure you can even vectorize my code to get the best of both worlds.