split list into sublists

Question

Define a function split_list that takes in a list of numbers and a number a and splits it into two sublists: list1 that contains numbers smaller than or equal to a, and another list, list2 containing numbers greater than a. list1 and list2 must be returned as elements of a tuple.

My code:

def split_list(lst, a):
    list1 = []
    list2 = []
    for i in lst:
        if i <= a:
            list1.append(i)
            lst.remove(i)
        elif i > a:
            list2.append(i)
            lst.remove(i)
    return (list1, list2)

Test code:

split_list([1, 10, 4, 9, 7, 2, 5, 8, 3, 4, 9, 6, 2], 5)

should give: ([1, 4, 2, 5, 3, 4, 2], [10, 9, 7, 8, 9, 6])

But I got ([1, 4, 5, 3, 2], [7, 9]) instead. Whats wrong with my code?

Answer 1

As commented by Martijn, you're breaking your list lst by trying to iterate through it and remove from it at the same time.

You don't need to remove from lst, so you can remove those lines.

Or, you could iterate over a copy of lst, whilst removing from the original lst. To do this, simply change your for loop to be:

for i in lst[:]

This will copy a slice corresponding to the whole of lst for you to iterate over.

Answer 2

To avoid skipping elements you can alternatively loop through the list backwards deleting as you go. Here's an example using list indexes instead of the "for listitem in list" construct.

def split_list(lst, a):
    list1 = []
    list2 = []
    for x in range(len(lst)-1,-1,-1):
        if lst[x] <= a:
            list1.append(lst[x])
            del lst[x]
        elif lst[x] > a:
            list2.append(lst[x])
            del lst[x]
    return (list1, list2)

print split_list([1, 10, 4, 9, 7, 2, 5, 8, 3, 4, 9, 6, 2], 5)    

# outputs ([2, 4, 3, 5, 2, 4, 1], [6, 9, 8, 7, 9, 10])

Answer 3

Modifying a list while iterating over it is always going to end badly. I would do instead:

for i in lst[:]:  # Make a shallow copy.
    #stuff

We are then iterating over that unmodified copy, so based on this we can edit the original safely.

Or just don't modify the list at all, meaning you can remove all uses of lst.remove(i) which should also fix the problem.

Answer 4

def split_list(lst, a):
    list1 = []
    list2 = []
    for i in lst:
        if i <= a:
            list1.append(i)

        elif i > a:
            list2.append(i)

    return (list1, list2)

There's no need to remove the items from the lst list as you iterate through it.