python - Iterating over numbers including prime factors

Question

I want to iterate over all numbers from 1 to n, and I need to know the primes factors of each number. A simple implementation would be to simply iterate over all numbers from 1 to n and calculate the prime factors of each, taking a time of O(n^(3/2)).

I've come up with this solution which generates each number exactly once, and one number (which is not yielded) extra per prime number, having a time complexity of O(n + p).

# primes contains enough primes so accessing it will never fail.
# Generate numbers 1 to limit (inclusive) including prime factors.
def gen(limit, i=0):
    # Primes above the limit are used 0 times, p**0 = 1
    if primes[i] > limit:
        yield 1, {}
        return

    # Generate all numbers using only higher primes...
    for n, f in gen(limit, i + 1):
        # ...and then add the current prime as often as possible each time.
        while n <= limit:
            yield n, f.copy()

            n *= primes[i]
            cf = f.get(primes[i], 0)
            f[primes[i]] = cf + 1

This prints the correct result for small limits. However, this solution is recursive, and reaches the maximum stack size quite fast. The output is as

1 {}
2 {2: 1}
4 {2: 2}
8 {2: 3}
3 {3: 1}
6 {3: 1, 2: 1}
9 {3: 2}
5 {5: 1}
10 {5: 1, 2: 1}
7 {7: 1}

Note that I don't expect the numbers to be generated in-order.

Is it possible to transform this into an iterative solution, or is there a simpler one entirely?

Answer 1

There is a simpler one using memoization. For example:

from functools import cache

@cache
def prime_factors(num):
    '''
    Example: 24 -> defaultdict({2: 3, 3: 1})
    '''
    assert isinstance(num, int) and num > 0

    # handle edge cases like 1
    if num == 1:
        return defaultdict(int)

    # find smallest divisor up to sqrt(num)
    # if no divisor found, num is prime
    # this can be optimized to only try prime divisors
    divisor, prime = 2, True
    while divisor**2 <= num:
        if num % divisor == 0:
            prime = False
            break
        divisor += 1

    if prime:
        factors = defaultdict(int)
        factors[num] += 1
    else:
        quotient = num // divisor
        # be sure to copy as dicts are mutable
        factors = prime_factors(quotient).copy()
        factors[divisor] += 1
    return factors

for i in range(10000):
    print(prime_factors(i))

This is still technically recursive, but in the memoized case you just do a lookup.