Python: speed up removal of every n-th element from list.

Question

I'm trying to solve this programming riddle and although the solution (see code below) works correctly, it is too slow for succesful submission.

• Any pointers as how to make this run faster (removal of every n-th element from a list)?
• Or suggestions for a better algorithm to calculate the same; seems I can't think of anything else than brute-force for now...

Basically, the task at hand is:

GIVEN:
L = [2,3,4,5,6,7,8,9,10,11,........]
1. Take the first remaining item in list L (in the general case 'n'). Move it to 
   the 'lucky number list'. Then drop every 'n-th' item from the list.
2. Repeat 1

TASK:
Calculate the n-th number from the 'lucky number list' ( 1 <= n <= 3000)

My original code (it calculated the 3000 first lucky numbers in about a second on my machine - unfortunately too slow):

"""
SPOJ Problem Set (classical) 1798. Assistance Required
URL: http://www.spoj.pl/problems/ASSIST/
"""

sieve = range(3, 33900, 2)
luckynumbers = [2]

while True:
    wanted_n = input()
    if wanted_n == 0:
        break

    while len(luckynumbers) < wanted_n:
        item = sieve[0]
        luckynumbers.append(item)
        items_to_delete = set(sieve[::item])
        sieve = filter(lambda x: x not in items_to_delete, sieve)
    print luckynumbers[wanted_n-1]

EDIT: thanks to the terrific contributions of Mark Dickinson, Steve Jessop and gnibbler, I got at the following, which is quite a whole lot faster than my original code (and succesfully got submitted at http://www.spoj.pl with 0.58 seconds!)...

sieve = range(3, 33810, 2)
luckynumbers = [2]

while len(luckynumbers) < 3000:
    if len(sieve) < sieve[0]:
        luckynumbers.extend(sieve)
        break
    luckynumbers.append(sieve[0])
    del sieve[::sieve[0]]

while True:
    wanted_n = input()
    if wanted_n == 0:
        break
    else:
        print luckynumbers[wanted_n-1]

Answer 1

You're better off using an array and zeroing out every Nth item using that strategy; after you do this a few times in a row, the updates start getting tricky so you'd want to re-form the array. This should improve the speed by at least a factor of 10. Do you need vastly better than that?

Answer 2

Try using these two lines for the deletion and filtering, instead of what you have; filter(None, ...) runs considerably faster than the filter(lambda ...).

sieve[::item] = [0]*-(-len(sieve)//item)
sieve = filter(None, sieve)

Edit: much better to simply use del sieve[::item]; see gnibbler's solution.

You might also be able to find a better termination condition for the while loop: for example, if the first remaining item in the sieve is i then the first i elements of the sieve will become the next i lucky numbers; so if len(luckynumbers) + sieve[0] >= wanted_n you should already have computed the number you need---you just need to figure out where in sieve it is so that you can extract it.

On my machine, the following version of your inner loop runs around 15 times faster than your original for finding the 3000th lucky number:

while len(luckynumbers) + sieve[0] < wanted_n:
    item = sieve[0]
    luckynumbers.append(item)
    sieve[::item] = [0]*-(-len(sieve)//item)
    sieve = filter(None, sieve)
print (luckynumbers + sieve)[wanted_n-1]

Answer 3

An explanation on how to solve this problem can be found here. (The problem I linked to asks for more, but the main step in that problem is the same as the one you're trying to solve.) The site I linked to also contains a sample solution in C++.

The set of numbers can be represented in a binary tree, which supports the following operations:

• Return the nth element
• Erase the nth element

These operations can be implemented to run in O(log n) time, where n is the number of nodes in the tree.

To build the tree, you can either make a custom routine that builds the tree from a given array of elements, or implement an insert operation (make sure to keep the tree balanced).

Each node in the tree need the following information:

• Pointers to the left and right children
• How many items there are in the left and right subtrees

With such a structure in place, solving the rest of the problem should be fairly straightforward.

I also recommend calculating the answers for all possible input values before reading any input, instead of calculating the answer for each input line.

A Java implementation of the above algorithm gets accepted in 0.68 seconds at the website you linked.

(Sorry for not providing any Python-specific help, but hopefully the algorithm outlined above will be fast enough.)

Answer 4

Why not just create a new list?

L = [x for (i, x) in enumerate(L) if i % n]

Answer 5

This series is called ludic numbers

__delslice__ should be faster than __setslice__+filter

>>> L=[2,3,4,5,6,7,8,9,10,11,12]
>>> lucky=[]
>>> lucky.append(L[0])
>>> del L[::L[0]]
>>> L
[3, 5, 7, 9, 11]
>>> lucky.append(L[0])
>>> del L[::L[0]]
>>> L
[5, 7, 11]

So the loop becomes.

while len(luckynumbers) < 3000:
    item = sieve[0]
    luckynumbers.append(item)
    del sieve[::item] 

Which runs in less than 0.1 second