pythonpython-3.xlist-comprehensionfactorization

Prime factorization using list comprehension in Python


How to write a function which returns a list of tuples like (ci, ki) for n such that n = c1k1 c2k2 ... ciki is a prime number?

For example: 12600 = 23 · 32 · 52 · 71

Desired output: [(2, 3), (3, 2), (5, 2), (7, 1)]

I know how to do it with while but is it possible to do it using list comprehension? Efficiency is not required in this task.

# naive function 
def is_prime(n):
    return n > 1 and all(n % i != 0 for i in range(2, n))

# while version
def factorization_while(n):
    res = []
    for i in range(1, n + 1):
        if is_prime(i):
            j = 0
            while n % i == 0:
                n = n // i
                j += 1
            if j != 0:
                res.append((i, j))
    return res

# list comprehension version
def factorization(n):
    return [
        (i, j) for i in range(1, n + 1) if is_prime(i) \
        and n % i == 0 ... # add smth
    ]

Solution

  • I don't think this should be too hard. You don't actually need to modify n to find its prime factors, they're all completely independent of each other. So just iterate through the appropriate primes, and find the maximum power that works!

    from math import log
    
    def prime_factors(n):
        return [(prime, max(power for power in range(1, int(log(n, prime))+1)
                                  if n % prime**power == 0))
                for prime in range(2, n+1) if n % prime == 0 and isprime(prime)]
    

    There are a few ways you might improve this further. You could use itertools.takewhile on an infinite generator of powers (e.g. itertools.count), as once you find the first power such that prime**power is not a factor of n, none of the later ones will be either. That would allow you to avoid the log call.

    And there are a whole bunch of ways to efficiently iterate over the primes (see for instance, the very simple generator suggested here, or higher performance versions that you can find in the answers to a few different questions here on SO).