This is Dijkstra's algorithm! He's super cool:

%%time
import math as MTX

\# scenario 1 : inputting new valid numbers
\# POOL = \{

\#     2   :   6,

\#     3   :   9,

\#     #should add 5

\# \}
\# scenario 2 : starting fresh

POOL = \{\}
\# scenario X : satisfying for more than one keys
\# POOL = \{

\#     2   :   12 ,

\#     3   :   12 ,

\#     5   :   25 ,

\#     7   :   49 ,

\#     11  :   121 ,

\#     13  :   169 ,

\# \}




def search_shallow(x):

count = 0

barrier = MTX.log(x)

valid_keys = \[\]
for i,j in POOL.items():

if count == barrier:     # experimental, unsure if it is necessary for it to run for such a limited amount

\#print(count)         # it is not obvious that it should be limited, however, it may work

return valid_keys      #update, so far it works
if x == j:

\#print(f"key : \{i\}")

valid_keys.append(i)

count += 1

return valid_keys

def increment(key_list):

for key in key_list:

POOL\[key\] += key
def extend(x):

POOL\[x\] = x**2

\# BODY

\# modify only here

\# \_\_____________\_\_

n = 1000000

\# \_\_____________\_\_

for n in range(2,n+1):

key_list = search_shallow(n)
\#print(key_list)

if len(key_list) \!= 0:

increment(key_list)

else:

extend(n)
print(POOL.keys())

print(len(POOL.keys()))