Count numbers that are the sum of three perfect cubes

Question

From numbers from 1 to 100, count the numbers where the following equation is satisfied with 1 <= a <= 100, 1 <= b <= 100, and 1 <= c <= 100: $$a³ + b³ + c³$$

I have this:

def has_power3(n):
    for i in range(1, n+1):
        for j in range(1, n+1):
            for k in range(1, n+1):
                if i**3 + j**3 + k**3 == n:
                    return True

    return False


powers_3 = 0
for num in range(1, 100+1):
    if has_power3(num):
        powers_3 += 1

print(powers_3)

However, this code above is O(n⁴). Is there a way to speed this code up? Any other suggestions?

Answer 1

Optimization: Code movement

You are repeatedly doing the same computations over and over. How many times is \$i^3\$ computed?

        for j in range(1, n+1):
            for k in range(1, n+1):
                if i**3 + j**3 + k**3 == n:
                    return True

Inside these loops, the value of i does not change. With j and k ranging over 100 values each, that is 10,000 redundant calculations! Since \$i^3\$ is constant, you can move that out of the loop. Ditto for \$j^3\$ and the middle loop:

    n_plus_1 = n + 1
    for i in range(1, n_plus_1):
        i3 = i ** 3
        for j in range(1, n_plus_1):
            i3_plus_j3 = i3 + j ** 3
            for k in range(1, n_plus_1):
                if i3_plus_j3 + k ** 3 == n:
                    return True

Optimization: Search Space

You're looking for

$$i^3 + j^3 + k^3 = n, 1 \le i \le 100, 1 \le j \le 100, 1 \le k \le 100$$

If, after you've exhausted searching \$i = 1\$, and not finding any values of j or k which work, is there any point in exploring \$j = 1\$ or \$k = 1\$ for \$i \gt 1\$? No! That would just be a permutation of a triplet you've already tried.

Without loss of generality, you can search a much smaller space:

$$i^3 + j^3 + k^3 = n, 1 \le i \le j \le k \le 100$$

As in:

    n_plus_1 = n + 1
    for i in range(1, n_plus_1):
        i3 = i ** 3
        for j in range(i, n_plus_1):
            i3_plus_j3 = i3 + j ** 3
            for k in range(j, n_plus_1):
                if i3_plus_j3 + k ** 3 == n:
                    return True

Optimization: Early termination

If \$i^3 + j^3 + k^3 \gt n\$, then trying with a larger value of k won't result in a smaller value, so you can stop the inner loop at that point.

Similarly, if \$i^3 + j^3 + j^3 \gt n\$, then there is no point in even entering the inner loop, and larger values of j are also pointless, so you can stop the middle loop at that point.

Finally, if \$i^3 + i^3 + i^3 \gt n\$, then there is no point in even entering the middle loop, and larger values of i are also pointless, so you can stop the outer loop.

    n_plus_1 = n + 1
    for i in range(1, n_plus_1):
        i3 = i ** 3
        if 3 * i3 > n:
            break
        for j in range(i, n_plus_1):
            j3 = j ** 3
            if i3 + 2 * j3 > n:
                break
            i3_plus_j3 = i3 + j3
            for k in range(j, n_plus_1):
                i3_plus_j3_plus_k3 = i3_plus_j3 + k ** 3
                if i3_plus_j3_plus_k3 > n:
                    break
                if i3_plus_j3_plus_k3 == n:
                    return True

Optimization: Loop end-points

The above added a lot of if statements to the search, complicating the algorithm. We can do math to determine the actual endpoints, and remove the if conditions from inside the loops. The inner loop completely. vanishes.

$$ 3 * i^3 > n \rightarrow i_{max} = \lfloor \sqrt[3]{n / 3} \rfloor $$

$$ i^3 + 2 * j^3 > n \rightarrow j_{max} = \lfloor \sqrt[3]{\frac{n - i^3} {2}} \rfloor $$

$$ i^3 + j^3 + k^3 = n \rightarrow k = \sqrt[3]{n - i^3 - j^3} $$

    THIRD = 1 / 3

    i_max = int((n / 3) ** THIRD)
    for i in range(1, i_max + 1):
        i3 = i ** 3

        j_max = int(((n - i3) / 2) ** THIRD)
        for j in range(i, j_max + 1):
            j3 = j ** 3

            k = int((n - i3 - j3) ** THIRD)
            if i3 + j3 + k ** 3 == n:
                return True

Optimization: Precomputed Values

As vnp mentions, we can simplify/remove the math in the last step by precomputing cubes. Storing them in a set allows for \$O(1)\$ lookup.

    THIRD = 1 / 3

    k_max = int((n - 2) ** THIRD)
    k_cubes = {k ** 3 for k in range(1, k_max + 1)}

    i_max = int((n / 3) ** THIRD)
    for i in range(1, i_max + 1):
        i3 = i ** 3

        j_max = int(((n - i3) / 2) ** THIRD)
        if any((n - i3 - j ** 3) in k_cubes for j in range(i, j_max + 1)):
            return True

---

Alternate Approach

Instead of asking "if 1 is the sum of 3 cubes", and then "if 2 is the sum of 3 cubes", and then "if 3 is the sum of 3 cubes" ... and so on up to "if 100 is the sum of 3 cubes", turn the problem around, and record a set of the sums of 3 cubes. The length of that set is the answer.

def cube_sums(n: int) -> set[int]:
    i_max = int((n / 3) ** (1 / 3)) + 1
    j_max = int(((n - 1) / 2) ** (1 / 3)) + 1
    k_max = int((n - 2) ** (1 / 3)) + 1

    return {t for i in range(1, i_max)
              for j in range(i, j_max)
              for k in range(j, k_max)
              if (t := i ** 3 + j ** 3 + k ** 3) <= n}

cubes = cube_sums(100)
print(len(cubes))

Answer 2

It doesn't make sense to exercise the entire range(1, n+1) for i. As soon as i reaches the cubic root of n, you may safely stop. Similarly, j shall only run until \$\sqrt[3]{n - i^3}\$.

Finally, you only need to determine if \$n - i^3 - j^3\$ is a perfect cube. For that, you may precompute a list of perfect cubes and binary search it, or use Newton-Rafson. Which one is faster, only profiling may tell.

Answer 3

All the answers are good, I just wanted to add a simple way to make the optimization process smoother. Keep the old inefficient function and make an "oracle test" comparing the old results with the new results, something like:

def oracle_test(slow_func, fast_func, inputs):
    for i in inputs:
        t = time.perf_counter()
        out_slow = slow_func(i)
        print(f"For input {i} slow took {time.perf_counter() - t}")
        
        t = time.perf_counter()
        out_fast = fast_func(i)
        print(f"For input {i} fast took {time.perf_counter() - t}\n")
        
        assert out_slow == out_fast
    return True

Now you can keep the very simple and obviously correct function as an "oracle" to check if the optimized version is still working correctly, so applying to your code, applying some optimizations suggested in other answers:

import time

def has_power3_fast(n):
    for i in range(1, n+1):
        i_3 = i**3
        if i_3 > n:
            continue
        for j in range(1, n+1):
            j_3 = j**3
            if i_3 + j_3 > n:
                continue
            for k in range(1, n+1):
                if i_3 + j_3 + k**3 == n:
                    return True

    return False

def has_power3_slow(n):
    for i in range(1, n+1):
        for j in range(1, n+1):
            for k in range(1, n+1):
                if i**3 + j**3 + k**3 == n:
                    return True

    return False

def oracle_test(slow_func, fast_func, inputs):
    for i in inputs:
        t = time.perf_counter()
        out_slow = slow_func(i)
        print(f"For input {i} slow took {time.perf_counter() - t}")
        
        t = time.perf_counter()
        out_fast = fast_func(i)
        print(f"For input {i} fast took {time.perf_counter() - t}\n")
        
        assert out_slow == out_fast
    return True

print(oracle_test(
    has_power3_slow,
    has_power3_fast,
    [50, 100, 200]
    )
)

In this way you can automatically test and benchmark all changes to greatly speed-up development of an optimized version.

For input 50 slow took 0.07776241900137393
For input 50 fast took 0.00014601400471292436

For input 100 slow took 0.5823429270021734
For input 100 fast took 0.00048202899779425934

For input 200 slow took 4.537595040004817
For input 200 fast took 0.001289013001951389

True

Answer 4

You can quite easily do it in O(n) with n units of storage.

First you should see that for example to check whether 100 is the sum of three perfect cubes, it is totally pointless to compute any of the cubes from 5^3 = 125 to 101^3, since adding two more cubes well give a result much bigger than 100. The number of iterations for each loop would be $n^{1/3}$, and the total iterations for three nested loops would be n. You examine n numbers, so that simple change gets you to $n^2$. You could have a lookup table that tells you if x is a cube and use that to determine whether a k exists, that would get you down to $n^{5/3}$ operations.

Instead realise that for every “num” you do mostly the same calculation and that is pointless. And we can assume that i <= j <= k since the order doesn’t matter. We must have 3i^3 <= n, then i^3 + 2j^3 <+ n, then i^3 +j^3 2 + j^3 <= n. So to check:

Create an array “isCube” containing indices 1 to n and fill them all with a value “false”. 
Loop for I from 1 until 3i^3 > n. 
    Loop for j from I until i^3+2j^3 > n
        Loop for k from j until i^3 + j^3+k^3 > n
            Set element i^3 + j^3 + k^3 to true. 

 Print all numbers if a bit in your array is set.