4
\$\begingroup\$

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?

\$\endgroup\$
1
  • \$\begingroup\$ I wonder if you could turn this problem into one that uses a graph algorithm, and if that would help \$\endgroup\$
    – moonman239
    Jan 28 at 0:39

4 Answers 4

12
\$\begingroup\$

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))
\$\endgroup\$
2
  • 3
    \$\begingroup\$ I think the cube roots in the final version are a mistake - we must be required to use the sums up to i = j = k = 100, not just to i³ + j³ + k³ = 100. At least, that's my reading of the question, since otherwise there would be no need to specify upper limits for i, j and k, which would naturally all be less than 5. But I see why you can read it the other way ("For numbers up to 100" is pretty ambiguous if you ask me). It's an excellent answer, and covers everything I was going to say. \$\endgroup\$ Jan 27 at 8:11
  • 1
    \$\begingroup\$ @Toby The OP was only testing for num in range(1, 100+1), which mathematically puts cube-root limits on the i,j,k ... but you're right, I'm totally ignoring the problem limits themselves. cube_sums(10_000_000) would give the wrong answer since it would include \$101^3+102^3+103^3\$ and so on. The problem description and the code are not consistent with each other ... or at least imply a non-listed constraint. \$\endgroup\$
    – AJNeufeld
    Jan 27 at 16:05
4
\$\begingroup\$

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.

\$\endgroup\$
2
\$\begingroup\$

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
\$\endgroup\$
0
1
\$\begingroup\$

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. 
\$\endgroup\$
1
  • 2
    \$\begingroup\$ You have presented an alternative solution, but haven't reviewed the code. Please edit to show what aspects of the question code prompted you to write this version, and in what ways it's an improvement over the original. It may be worth (re-)reading How to Answer. \$\endgroup\$ Jan 27 at 13:25

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.