python get first x elements of a list and remove them-CodePudding

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?

CodePudding user response：

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 list[start:end]
    return partial

first_list = splice(lst, end = 2)

CodePudding user response：

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]

CodePudding user response：

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]

CodePudding user response：

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])