astroboy
astroboy

Reputation: 1078

python get first x elements of a list and remove them

Note: I know there is probably an answer for this on StackOverflow already, I just can't find it.

I need to do this:

>>> lst = [1, 2, 3, 4, 5, 6]
>>> first_two = lst.magic_pop(2)
>>> first_two
[1, 2]
>>> lst
[3, 4, 5, 6]

Now magic_pop doesn't exist, I used it just to show an example of what I need. Is there a method like magic_pop that would help me to do everything in a pythonic way?

Upvotes: 6

Views: 906

Answers (6)

no comment
no comment

Reputation: 10465

Slice and del-slice as Barmar showed is really the straightforward and best, but here's a more fun alternative:

first_x, lst[:x] = lst[:x], ()

Like Barmar's, it takes only linear time and literally removes the elements from the given list as the question asks for. (None of the other posted solutions have these two qualities, although in many cases, lst = lst[x:] works for "removing", albeit less efficiently.)

Little benchmark, getting and removing/"removing" the first thousand of a million elements:

  0.7 ms  first_x = lst[:x]; del lst[:x]
  0.7 ms  first_x, lst[:x] = lst[:x], ()
  0.7 ms  first_x = lst[:x]; lst[:x] = ()
 31.1 ms  first_x = lst[:x]; lst = lst[x:]
1000.8 ms  first_x = list(map(lst.pop, [0] * x))

Test/benchmark code:

from timeit import timeit

codes = [
    'first_x = lst[:x]; del lst[:x]',
    'first_x, lst[:x] = lst[:x], ()',
    'first_x = lst[:x]; lst[:x] = ()',
    'first_x = lst[:x]; lst = lst[x:]',
    'first_x = list(map(lst.pop, [0] * x))',
]

for code in codes:
    lst = [1, 2, 3, 4, 5, 6]
    x = 2
    exec(code)
    print(first_x, lst)

setup = '''
lst = list(range(1000000))
x = 1000
'''
for code in codes:
    t = timeit(code, setup, number=1)
    print(f'{t*1e3:5.1f} ms ', code)

Attempt This Online!

Upvotes: 1

Ishaan Rao
Ishaan Rao

Reputation: 11

A very simple way to do this using simple python code would be to simply index the first 2 elements of the list, copy them to a second empty "container" list, and then remove the same 2 items:

sample_list = ["a","b","c","d","e","f"] # Original list
new_list = [] # Container list for first 2 items

for i in range(2): # First 2 indices
    new_list += sample_list[0] # Append item at index 0 -> new list
    sample_list.pop(0) # Remove item at index 0
    
print(sample_list) # --> ['c','d','e','f']
print(new_list) # --> ['a','b']

Output:

['c','d','e','f'] 
['a','b']

Upvotes: 0

Barmar
Barmar

Reputation: 782166

Do it in two steps. Use a slice to get the first two elements, then remove that slice from the list.

first_list = lst[:2]
del lst[:2]

If you want a one-liner, you can wrap it in a function.

def splice(lst, start = 0, end = None):
    if end is None:
        end = len(lst)
    partial = lst[start:end]
    del lst[start:end]
    return partial

first_list = splice(lst, end = 2)

Upvotes: 7

niko
niko

Reputation: 5281

One option would be using slices

def func(lst: list, n: int) -> tuple:
    return lst[:n], lst[n:]


lst = [1, 2, 3, 4, 5, 6]
first_two, lst = func(lst, 2)
print(lst)
# [3, 4, 5, 6]

Alternative, here is a kafkaesque approach using builtins:

def func(lst: list, i: int) -> list:
    return list(map(lambda _: lst.pop(0), range(i)))

lst = list(range(10))
first_two = func(lst, 2)
print(first_two)
# [0, 1]
print(lst)
# [2, 3, 4, 5, 6, 7, 8, 9]

Upvotes: 3

Anurag Reddy
Anurag Reddy

Reputation: 1225

You can just create a lambda function to do this.

magic_pop = lambda x,y:(x[:y],x[y:]) if y<len(lst) else x
magic_pop(lst,3)
Out[8]: ([1, 2, 3], [4, 5, 6])

Upvotes: 0

Andrej Kesely
Andrej Kesely

Reputation: 195573

Try:

lst = [1, 2, 3, 4, 5, 6]

first_two, lst = lst[:2], lst[2:]

print(first_two)
print(lst)

Prints:

[1, 2]
[3, 4, 5, 6]

Upvotes: 5

Related Questions