Abhinav Ramakrishnan
Abhinav Ramakrishnan

Reputation: 1090

list comprehension with multiple assignments

I presently have this code for factoring large numbers:

def f1(n):
    return [[i, n//i] for i in range(1 , int(n**0.5) + 1) if n % i == 0]

It's the fastest version I've seen so far (If there's a faster way I'd love to know about that as well), but I'd like a single list of all the factors with no nesting (so I want something like: [factor 1, factor 2, factor 3,..., factor n-3, factor n-2, factor n-1, factor n] and so on. The order isn't really important.

As such I was wondering if there was a way to ascribe multiple assignments via a list comprehension.

i.e.

def f1(n):
    return [i, n//i for i in range(1 , int(n**0.5) + 1) if n % i == 0]

That way I don't have a nested list. It would be faster and speed is of the essence.

I looked in the documentation and I couldn't find a single example of multiple assignments.

Upvotes: 1

Views: 1340

Answers (3)

Padraic Cunningham
Padraic Cunningham

Reputation: 180482

Use itertools.chain:

from itertools import chain

def f1(n):
    return list(chain.from_iterable([i, n//i] for i in xrange(1 , int(n**0.5) + 1) if not n % i))

If you don't need a list remove the list call on chain and just iterate over the returned chain object.

If optimization is important you should use extend and xrange:

def f1(n):
    l = []
    for i in xrange(1, int(n**0.5)+1):
        if not n % i:
            l.extend((i,n//i))
    return l

Upvotes: 1

Dunes
Dunes

Reputation: 40833

List comprehensions are great, but sometimes they're not best the solution, depending on requirements for readability and speed. Sometimes, just writing out the implied for loop (and if statement) is more readable and quicker.

def factors(n):
    l = []
    for i in range(1, int(n**0.5)+1):
        if n % i == 0:
            l.append(i)
            l.append(n//i)
    return l

For small numbers, the above function is quicker than the list comprehension. At larger numbers (1,000,000 and bigger), the function and list comprehension are equal in terms of speed.

For a slight speed increase you can also cache the append method of the list, though this makes the function slightly less readable.

def factors(n):
    l = []
    append = l.append
    for i in range(1, int(n**0.5)+1):
        if n % i == 0:
            append(i)
            append(n//i)
    return l

Speed comparison:

In [86]: %timeit factors_list_comprehension(1000)
100000 loops, best of 3: 7.57 µs per loop

In [87]: %timeit factors_function(1000)
100000 loops, best of 3: 6.24 µs per loop

In [88]: %timeit factors_optimised_function(1000)
100000 loops, best of 3: 5.81 µs per loop

In [89]: %timeit factors_list_comprehension(1000000)
10000 loops, best of 3: 111 µs per loop

In [90]: %timeit factors_function(1000000)
10000 loops, best of 3: 108 µs per loop

In [91]: %timeit factors_optimised_function(1000000)
10000 loops, best of 3: 106 µs per loop

Upvotes: 2

Nayuki
Nayuki

Reputation: 18552

You can achieve the desired result using sum(). For example:

>>> sum([[1,6],[2,3]],[])
[1, 6, 2, 3]

We can define the answer in terms of your existing code:

def f2(n):
    return sum(f1(n), [])

However, be careful that your code returns the square root twice when n is a perfect square:

>>> f1(9)
[[1, 9], [3, 3]]
>>> f2(9)
[1, 9, 3, 3]

Upvotes: 0

Related Questions