Optimization of python code for calculating list of squared divisors

Question

I was participating in a python challenge in codewars website. I encountered the following challenge:

Divisors of 42 are : 1, 2, 3, 6, 7, 14, 21, 42. These divisors squared are: 1, 4, 9, 36, 49, 196, 441, 1764. The sum of the squared divisors is 2500 which is 50 * 50, a square!

Given two integers m, n (1 <= m <= n) we want to find all integers between m and n whose sum of squared divisors is itself a square. 42 is such a number.

The result will be an array of arrays, each subarray having two elements, first the number whose squared divisors is a square and then the sum of the squared divisors.

The output should be:

list_squared(1, 250) --> [[1, 1], [42, 2500], [246, 84100]]
list_squared(42, 250) --> [[42, 2500], [246, 84100]]
list_squared(250, 500) --> [[287, 84100]]

I have written following code with two additional functions: one corresponding to determine all factors of a number and other checking if a number is perfect square or not.

Function to determine all factors:

def fact(m):
    return [i for i in range(1,m+1) if m%i == 0]

Function to check if a number is perfect square and return 0 if it is not otherwise return square root

def square_root(x):
    ans = 0
    while ans < x // 2 + 1:
        ans = ans + 1

        if ans*ans == x:
            return ans
            break;
    return 0

Function where the desired result is calculated

def list_squared(m, n):
    # your code
    fac=[]
    for i in range(m,n+1):
        sq_sum = sum(list(map(lambda x: x**2,fact(i))))
        if square_root(sq_sum) != 0:
            fac.append([i,sq_sum])
    return fac

This code gives me the correct result, however it is too long. I was able to pass all the test results but it took me around 6000 ms. When I attempted to submit the code, the web submission returns that the algorithm is inefficient and it took more than 1200 ms which is the maximum.

I would highly appreciate if anyone can point to a better algorithm for this.

Answer 1

There are several optimizations to your code but the biggest one is to stop when ans*ans becomes bigger than x here:

def square_root(x):
    ans = 0
    while True:
        ans += 1
        sqans = ans*ans
        if sqans == x:
            return ans
        elif sqans > x:
            return 0

The condition in the while can be removed, since now the test is done on the square value.

with that optimization, I drop from 8 seconds to 0.07 seconds with the 250, 500 case.

But that's stil not satisfactory. In general, algorithms containing a condition to break or return are at least O(n) and even if you can save time, the complexity is too high.

You can do better by simply checking the square of the rounded square root:

def square_root(x):
    ans = int(x**0.5 + 0.5)  # rounded just in case it goes below the actual value (float inaccuracy)
    sqans = ans*ans
    return 0 if sqans !=x else x

I divide the execution time by further 2 with that (confirmed by Optimized way to find if a number is a perfect square)

Aside (that doesn't speed up that much but worth mentionning):

no need to convert map to list in sum:

sq_sum = sum(map(lambda x: x**2,fact(i)))

Also fact could avoid looping to the max number. Loop to the max number divided by 2 and add the max number to the list is equivalent. No more divisors exist above max number / 2

def fact(m):
    return [i for i in range(1,m//2+1) if m%i == 0] + [m]

Final edit: this is still slow because of the list comprehension used in fact. I could cut the time drastically by using a generator instead and add m*m outside it:

def sqfact(m):
    return (i*i for i in range(1,m//2+1) if m%i == 0)

final code, now runs so fast I get 0 seconds.

def sqfact(m):
    return (i*i for i in range(1,m//2+1) if m%i == 0)

def square_root(x):
    ans = int(x**0.5 + 0.5)
    return 0 if ans*ans !=x else x

def list_squared(m, n):
    # your code
    fac=[]
    for i in range(m,n+1):
        sq_sum = sum(sqfact(i)) + i*i  # add i square outside
        if square_root(sq_sum):
            fac.append([i,sq_sum])
    return fac

Answer 2

I have updated the fact function which was very inefficient. Now, rather than iterating to full value of m to find its factors, I am only going up to sqrt(m). This has reduced the run time immensely. The logic behind this is trivial so I am not elaborating. Following is the new code which worked for me.

def fact(m):

     #determining the lower factors i.e., smaller than sqrt(m)
    fac = [i for i in range(1, int(m**0.5) + 1) if m%i == 0]

     #determining the higher factors i.e., larger than sqrt(m)
    fac = fac + [m//i for i in range(1, int(m**0.5) + 1) if m%i == 0] 

    return sorted(list(set(fac))) #in order to get rid of duplicate factors