Vectorized PoC

import time
import numpy as np
tic = time.time()
valids = []
for digits in range(1,6):
    A = np.array(list(range(10**(digits-1), 10**digits)))
    Abase = A*10**digits
    Adiff = 1-2*A
    sqrt_delta = np.sqrt(Adiff**2 - 4*(A**2-Abase))
    B = (Adiff+sqrt_delta)//2
    num = Abase + B
    valid = num[(A+B)**2 == num].astype(np.int32)
    valids.append(valid)
toc = time.time()
print(len(np.hstack(valids)))
print(f"Time elapsed {toc-tic:.1f}s")