Four Divisors -- Leetcode 1390

Optimizing with number theory and the Sieve of Eratosthenes

leetcodemath

Today’s Daily Leetcode problem (01-04-2025) is Four Divisors.

This problem requires you to find all integers in an array with exactly four divisors and sum their divisors. Along the way we bump into some important math concepts and optimizations which make for a great showcase problem.

Overview

As a quick reminder: dividend = quotient * divisor + remainder.

For integer division, a number d is a divisor of another number num if num % d == 0. In other words, if dividing the dividend by the divisor produces an integer quotient and a remainder of 0.

Let’s look at the path of solutions:

  1. Brute force — iterate through arr, computing the number of divisors by dividing num by all d in range [1, num]. Add all divisors to a set to ensure no duplicates. Add the sum of the divisors to the total if there is only 4 of them.

  2. Optimized Brute force — search only in range [1, sqrt(num)]. When a divisor is found, we can get the quotient (ie the other divisor) using the formula from above: num/d.

    • This works because any divisor greater than sqrt(num) would have been found as the quotient of a smaller divisor.
  3. Number Theory — by recognizing the only 2 cases where a number has 4 divisors, we can quickly check for both cases using a primality test (which itself can be optimized).

Number Theory

References:

To optimize our solution, we need to think through what the problem is asking:

When does an integer have exactly four divisors?

Other than 0 and 1, all positive integers will have at least 2 divisors: 1 and itself. For primes, these two will be their only divisors. For non-primes, we can rely on the Fundamental theorem of arithmetic, noting that all integers will have a unique factorization of prime numbers.

This leads to two possible cases for having exactly four divisors:

  1. (1, p, q, p * q) where p, q are prime and p != q
  2. (1, p, p^2, p^3) where p is prime

For Case 1, p * q is n’s unique prime factorization. Since the number is the product of two primes, we know it will have no other divisors. We require p != q as this would only count as a single divisor.

For Case 2, the unique prime factorization is p * p * p. If we divide num by p we get p^2 which we know has no other divisors as it is the product of two primes.

This may seem a bit handwavey, but thinking through the first 10 numbers you will see that 6, 10 are examples of Case 1 and 8 is an example of Case 2. The Leetcode Editorial also has a good explanation for determining what the factor count for a given number is.

Case 1 in code:

def is_product_of_two_distinct_primes(num: int) -> tuple[bool, int, int]:
    """Check if num = p*q where p,q are distinct primes.
    Returns: (is_product, p, q) where is_product indicates if condition holds."""
    d = 2
    while d * d <= num:
        if num % d == 0:
            q = num // d
            if d == q:  # prime square only has 3 divisors
                return (False, 0, 0)
            if d != q and is_prime(d) and is_prime(q):
                return (True, d, q)
        d += 1
    return (False, 0, 0)

Case 2 in code:

def is_prime_cube(num: int) -> tuple[bool, int]:
    """Check if num = p^3 where p is prime.
    Returns: (is_cube, p) where is_cube indicates if condition holds."""
    p = int(round(num ** (1/3)))
    if p ** 3 == num and is_prime(p):
        return (True, p)
    return (False, 0)

Primality Test

The most general way to check if a number is prime is to check if any number smaller than it divides it (making a special exception for 0 and 1). This is called trial division:

def is_prime(n: int) -> bool:
    """Check if n is prime using trial division.
    Returns: True if n is prime, False otherwise."""
    if n < 2:
        return False
    for i in range(2, n):
        if n % i == 0:
            return False
    return True

A slightly more optimized way to do this would be to quickly check if n is even, and then limit the range of the divisor search to only odd numbers in the range [3, sqrt(n)] as any divisor > sqrt(n) would just be the quotient of a smaller divisor in the range we already tried.

def is_prime(n: int) -> bool:
    """Check if n is prime using optimized trial division.
    Returns: True if n is prime, False otherwise."""
    # 2 is prime, 0,1 and all other evens are composite
    if n == 2:
        return True
    if n < 2 or n % 2 == 0:
        return False

    # check if any odd num between [3, sqrt(n)] is a divisor
    d = 3
    while d * d <= n:  # ie, d <= n**0.5
        if n % d == 0:
            return False
        d += 2

    return True

There is LOTS of research in this space so there are many more optimizations and algorithms to be used, but this is good enough for a general use case.

At this point, we have everything we need for a solution. Two further optimizations to make would be to:

  1. Cache the result of any number we’ve already checked for four divisors (ie, in a hash map or using @functools.cache)
  2. Precomputing which numbers are prime (using the Sieve of Eratosthenes) instead of using a primality test for each num

Sieve of Eratosthenes

The Sieve of Eratosthenes seems scary at first, but it is rather simple.

The goal is to find all prime numbers up to some limit N. In our case, this is useful for precomputing primes so all calls to is_prime(n) return in constant time.

The main idea is that we want to filter (sieve) out the primes from the range of numbers from [2, N]. We do this by marking all multiples of each prime number in the range as non-prime, starting from 2. Any number we did not mark by the end of the range is a prime.

Here’s a concrete example sieving up to 10:

Initial: [2, 3, 4, 5, 6, 7, 8, 9, 10] all marked as potentially prime

i=2 (prime): Mark multiples 4, 6, 8, 10 as composite
i=3 (prime): Mark multiples 6, 9 as composite (6 already marked)
i=4: Already marked composite, skip
i=5 (prime): 5*5 = 25 > 10, so no multiples to mark

Result: [2, 3, 5, 7] remain marked as prime

Code:

def sieve(n: int) -> set[int]:
    """Returns set of all primes in the range [0, n].
    Returns: Set containing all prime numbers from 0 to n inclusive."""
    is_prime = [True for _ in range(n+1)]
    is_prime[0], is_prime[1] = False, False

    for i in range(2, n + 1):
        if not is_prime[i]:
            continue

        # Mark all multiples of i as composite
        # Start at i*i because smaller multiples were already marked by previous primes
        # Example: for i=5, we skip 5*2=10 (marked by 2) and 5*3=15 (marked by 3)
        for j in range(i*i, n + 1, i):
            is_prime[j] = False

    return {i for i in range(n + 1) if is_prime[i]}

A few simple ways to optimize would be to:

  1. Sieve only up until sqrt(N)
  2. Skip even multiples
  3. Start marking multiples from i*i as all previous multiples would have been counted by a previous prime

As above, there is a LOT of other optimizations and algorithms to use but this is everything we need for a near optimal solution.

Solution

Putting it all together:

def sieve(n: int) -> set[int]:
    """Returns set of all primes in the range [0, n].
    Returns: Set containing all prime numbers from 0 to n inclusive."""
    is_prime = [True for _ in range(n+1)]
    is_prime[0], is_prime[1] = False, False

    for i in range(2, int(n**0.5) + 1):
        if not is_prime[i]:
            continue

        # Start at i*i since smaller multiples already marked by previous primes
        for j in range(i * i, n + 1, i):
            is_prime[j] = False

    return {i for i in range(n + 1) if is_prime[i]}


class Solution:
    def sumFourDivisors(self, nums: list[int]) -> int:
        """Sum divisors for numbers with exactly four divisors.
        Returns: Sum of all four divisors for qualifying numbers in nums."""
        max_num = max(nums)
        prime_set = sieve(max_num)

        def is_prime(n: int) -> bool:
            """O(1) primality check using precomputed sieve.
            Returns: True if n is prime, False otherwise."""
            return n in prime_set

        def is_prime_cube(num: int) -> tuple[bool, int]:
            """Check if num = p^3 where p is prime.
            Returns: (is_cube, p) where is_cube indicates if condition holds."""
            p = int(round(num ** (1/3)))
            if p ** 3 == num and is_prime(p):
                return (True, p)
            return (False, 0)

        def is_product_of_two_distinct_primes(num: int) -> tuple[bool, int, int]:
            """Check if num = p*q where p,q are distinct primes.
            Returns: (is_product, p, q) where is_product indicates if condition holds."""
            d = 2
            while d * d <= num:
                if num % d == 0:
                    q = num // d
                    if d != q and is_prime(d) and is_prime(q):
                        return (True, d, q)
                d += 1
            return (False, 0, 0)

        cache = {}
        total = 0

        for num in nums:
            if num in cache:
                total += cache[num]
                continue

            divisor_sum = 0

            # Case 1: p^3 where p is prime
            is_cube, p = is_prime_cube(num)
            if is_cube:
                divisor_sum = 1 + p + p**2 + num
            # Case 2: p*q where p,q are distinct primes
            else:
                is_product, p, q = is_product_of_two_distinct_primes(num)
                if is_product:
                    divisor_sum = 1 + p + q + num

            cache[num] = divisor_sum
            total += divisor_sum

        return total

Complexity

For most intents and purposes, we can do a simplified analysis of the time and space complexity without diving too much into the exact asymptotic analysis. Just note that the actual bounds are a bit tighter than this, particularly the sieve implementation as O(M log log M).

Time Complexity: O(M log M + N·√M)

Where M = max(nums), N = len(nums), U = len(set(nums))

  • Sieve: O(M log M)
    • Outer loop: √M iterations
    • Inner loop: for each prime p, mark M/p multiples
    • Total work: M/2 + M/3 + M/5 + ... ≈ M log M
  • Per number: O(√M) worst case
    • is_prime_cube: O(1)
    • is_product_of_two_distinct_primes: O(√num) ≤ O(√M)
      • Loop runs up to √num
      • is_prime lookups are O(1)
  • Total: O(M log M) for sieve + O(N·√M) for processing all numbers
  • With caching duplicate nums: O(M log M + U·√M)

Space Complexity: O(M + N) or O(M) without cache

  • Sieve: O(M) - we store all primes up to M in the set
  • Cache (optional): O(U) ≤ O(N) for storing results of unique numbers
  • Total: O(M + N) with cache, O(M) without

Trade-off: Caching saves time on any duplicates we may encounter at the cost of O(N) extra space. If all nums are distinct or memory is tight, you can skip the caching.

Closing Notes

This is a fun problem with some fun theory to think through:

  • Why can we reduce our divisors search from [1, num] to only odd numbers in [3, sqrt(num)]?
  • How are we guaranteed that there are only two cases when a number has 4 divisors?
  • In the sieve, why do we start marking multiples from i*i? Why are previous multiples guaranteed to be marked?

We also can apply what we’ve learned to other Leetcode math problems like:

Some other fun resources I bumped into:

  1. The Genuine Sieve of Eratosthenes — writing an efficient functional implementation of the Sieve of Eratosthenes
  2. Algorithms for Competitive Programming: Sieve of Eratosthenes