Updating the list/set to be iterated in a for loop

Question

I'm trying to make a set of primes using Sieve's method. But it seems that prime = prime - set(range(2*n,N,n)) does not update the set prime in for n in prime:.

N = int(raw_input()) + 1
prime = set(range(2,N))

for n in prime:
    print n
    prime = prime - set(range(2*n,N,n))

For example, I input N = 100. The output is 1,2,3,...,100, making the code super slow. What should I do to update the set consistently?

Edit: A little deeper: Can you explain the difference between augmented and non-augmented assignments? Is there really no way I can accomplish with just one set? Because it just seems so naive and I can't quite believe such high-level language as Python cannot achieve this, since in C++ I can do it with just an array of boolean values

Answer 1

The for loop evaluates the iterable expression just once. You are not altering prime, you are rebinding the name to a new object (the output of prime - set(range(2*n,N,n)), the for loop will never see that.

Note also that a for loop over a set will do so in arbitrary order; the numbers are not sorted. You could easily be handed a non-prime first.

If you used augmented assignment to alter prime in-place, however, you'll get an error message about altering the set while iterating:

>>> N = 50
>>> prime  = set(range(2, N))
>>> for n in prime:
...     print n
...     prime -= set(range(2 * n, N, n))
... 
2
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Set changed size during iteration

You'll have to loop directly over the range() and track non-primes instead:

non_prime = set()
primes = []

for n in range(2, N):
    if n in non_prime:
        continue
    non_prime |= set(range(2 * n, N, n))
    primes.append(n)

Demo:

>>> non_prime = set()
>>> primes = []
>>> for n in range(2, N):
...     if n in non_prime:
...         continue
...     non_prime |= set(range(2 * n, N, n))
...     primes.append(n)
... 
>>> primes
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]

or you could use a while loop with a candidates set you empty:

candidates = set(range(2, N))
primes = []

while candidates:
    n = min(candidates)
    candidates -= set(range(2 * n, N, n))
    candidates.remove(n)
    primes.append(n)

but this requires looping over all of candidates each time to find the minimum value.

Note that you can further reduce the range of sizes by using range(n * n, N, n) here; all 2 * n values past 2 * 2 would already have been eliminated, all 3 * n past 3 * 3 as well, etc.:

>>> primes = []
>>> for n in range(2, N):
...     if n in non_prime:
...         continue
...     non_prime |= set(range(n * n, N, n))
...     primes.append(n)
... 
>>> primes
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]

Last but not least, the common way to implement the sieve is to use a list of booleans to mark out non-primes:

sieve = [True] * N
sieve[0] = sieve[1] = False
primes = []

for (n, is_prime) in enumerate(sieve):
    if is_prime:
        primes.append(n)
        for n in range(n * n, N, n):
            sieve[n] = False