Another short post, based on an interview question I was recently asked during a phone screen. The question was:

You are given a list of primes. Output a list of all possible products of them. For example, if given [2, 3, 11] output [1, 2, 3, 6, 11, 22, 33, 66].

Let's do the math first. Assuming that the list is unsorted and contains duplicates, our first goal would be to calculate how many times each prime appears in the list. This takes linear time. We will use Python's collections.Counter for this purpose.

Now, assume the list has $n$ distinct primes in it, denote the distinct primes by $p_i$ for $0\le{i}<n$ . Let $c_i$ count the number of times $p_i$ appears in the list. How many possible products are there? For each prime, we have $c_i+1$ choices for its exponent in the product. Hence the total number of products is given by

Our code will be based on this too. We have to go through all such possible choices for each prime. This is a classic problem in combinatorial generation, namely generating the Cartesian product of a finite collection of finite sets. Knuth discusses this in his latest volume of the Art of Computer Programming which you can read for a much more in-depth analysis of this problem.

First, let's do a simple solution based on a recursive idea. If the list of primes is empty then the resulting list should be [1]. That is our base case. We can then at each step recursively append all multiples of a new prime with the existing products, yielding the final list. While the idea is recursive, the implementation is simpler to do without a recursive call. Here it is:

def products(primes):
    result = [1]
    for p in primes:
        result += [x * p for x in result]
    return result

# Usage:
print products([2, 3, 11])
# Output:
# [1, 2, 3, 6, 11, 22, 33, 66]

The only issue here is that we need exponential memory to store the resulting list.

Can we do better? In terms of memory, yes. If instead of a recursion, we make use of Python's itertools module, we can only keep track of the current exponents of each prime, thereby using only linear memory. More specifically we need the itertools.product function which generates the Cartesian product of the iterators it is given. For more details, see my post on iterators and generators . Another advantage of this solution is that the result is a generator, meaning we do not need to keep the entire list in memory at once.

import itertools
import collections
import operator


def power_generator(p, n):
    """Generate all powers of p starting from 0 to n, inclusive."""
    q = 1
    for __ in xrange(n + 1):
        yield q
        q *= p


def products(primes):
    counts = collections.Counter(primes)
    p_generators = (power_generator(p, k) for p, k in counts.iteritems())
    for product in itertools.product(*p_generators):
        yield reduce(operator.mul, product, 1)


# Usage:
print list(products([2, 3, 11]))
print list(products([2, 2, 3, 5]))
print list(products([2, 2, 2, 2]))

# Output:
# [1, 3, 2, 6, 11, 33, 22, 66]
# [1, 5, 3, 15, 2, 10, 6, 30, 4, 20, 12, 60]
# [1, 2, 4, 8, 16]