Sirui Li
Sirui Li

Reputation: 265

How to extract elements from a nested list

I have a nested list in the form of [[1,2,3], [3,4,5], [8,6,2,5,6], [7,2,9]] I would like to extract every first item into a new list, every second item into a new list and the rest into a new nested list:

a = [1,3,8,7] b = [2,4,6,2], c = [[3], [5], [2,5,6],[9]]

Is it possible to avoid using the for loop because the real nested list is quite large? Any help would be appreciated.

Upvotes: 0

Views: 4851

Answers (4)

Toon Tran
Toon Tran

Reputation: 378

Suppose the nested list has unknown depth, then we'd have to use recursion

def get_elements(l):
    ret = []
    for elem in l:
        if type(elem) == list:
            ret.extend(get_elements(elem))
        else:
            ret.append(elem)
    return ret

l = [1,2,[3,4],[[5],[6]]]
print(get_elements(l)) 
# Output: [1, 2, 3, 4, 5, 6]

Though it is not quite recommended to use unknown-depth nested lists in the first place.

Upvotes: 0

Pitto
Pitto

Reputation: 8579

At the moment I cannot think a solution without for loops, I hope I will be able to update my answer later.

Here's a solution using for loops:

data = [[1,2,3], [3,4,5], [8,6,2,5,6], [7,2,9]]


list1 = []
list2 = []
list3 = []


for item in data:
    else_list = []
    for index, value in enumerate(item): 
        if index == 0:
            list1.append(value)
        elif index == 1:
            list2.append(value)
        else:
            else_list.append(value)
    list3.append(else_list)            
        
print(list1)
print(list2)
print(list3)

Output

[1, 3, 8, 7]
[2, 4, 6, 2]
[[3], [5], [2, 5, 6], [9]]

Just for fun I share also a performance comparison, great job in using just one for loop Meysam!

import timeit

# a = [1,3,8,7] b = [2,4,6,2], c = [[3], [5], [2,5,6],[9]]


def solution_1():
    data = [[1, 2, 3], [3, 4, 5], [8, 6, 2, 5, 6], [7, 2, 9]]

    list1 = []
    list2 = []
    list3 = []

    for item in data:
        else_list = []
        for index, value in enumerate(item):
            if index == 0:
                list1.append(value)
            elif index == 1:
                list2.append(value)
            else:
                else_list.append(value)
        list3.append(else_list)


def solution_2():
    arr = [[1, 2, 3], [3, 4, 5], [8, 6, 2, 5, 6], [7, 2, 9]]

    first_arr, second_arr, third_arr = [], [], []
    for nested in arr:
        first_arr.append(nested[0])
        second_arr.append(nested[1])
        third_arr.append(nested[2:])


def solution_3():
    l = [[1, 2, 3], [3, 4, 5], [8, 6, 2, 5, 6], [7, 2, 9]]
    a = [i[0] for i in l]
    b = [i[1] for i in l]
    c = [i[2:] for i in l]


if __name__ == "__main__":
    print("solution_1 performance:")
    print(timeit.timeit("solution_1()", "from __main__ import solution_1", number=10))
    print("solution_2 performance:")
    print(timeit.timeit("solution_2()", "from __main__ import solution_2", number=10))
    print("solution_3 performance:")
    print(timeit.timeit("solution_3()", "from __main__ import solution_3", number=10))

Output

solution_1 performance:
9.580000000000005e-05
solution_2 performance:
1.7200000000001936e-05
solution_3 performance:
1.7499999999996685e-05

Upvotes: 1

Tim Jim
Tim Jim

Reputation: 670

This is a naive, simple looped solution using list comprehensions, but see if it is fast enough for you.

l = [[1,2,3], [3,4,5], [8,6,2,5,6], [7,2,9]]
a = [i[0] for i in l]
b = [i[1] for i in l]
c = [i[2:] for i in l]

which returns:

>>a
[1, 3, 8, 7]

>>b
[2, 4, 6, 2]

>>c
[[3], [5], [2, 5, 6], [9]]

Upvotes: 2

Meysam
Meysam

Reputation: 718

Ultimately, whatever your solution would be, you're gonna have to have a for loop inside your code and my advice would be to make it as clean and as readable as possible.

That being said, here's what I would propose:

arr = [[1,2,3], [3,4,5], [8,6,2,5,6], [7,2,9]]

first_arr, second_arr, third_arr = [], [], []
for nested in arr:
    first_arr.append(nested[0])
    second_arr.append(nested[1])
    third_arr.append(nested[2:])

Upvotes: 4

Related Questions