Iterate through a nested list and pick certain elements and create a new list

Viewed 1264

An example:

list = [[2, 1, 2, 3, 4],
       [0, 4, 5],
       [1, 8, 9]]

So the first index inside a nested list decides which following numbers will be put into an unnested list.

[2, 1, 2, 3, 4] -> 2: so 1 and 2 gets picked up
[0, 4, 5] -> 0: no number gets picked up
[1, 8, 9] -> 1; number 8 gets picked up

Output would be:

[1, 2, 8]

This is what I have so far:

def nested_list(numbers):
    if isinstance(numbers[0], list):
        if numbers[0][0] > 0:
            nested_list(numbers[0][1:numbers[0][0] + 1])
    else:
        numbers = list(numbers[0])

    return numbers + nested_list(numbers[1:])

I try to get the list through recursion but something is wrong. What am I missing or could this be done even without recursion ?

5 Answers

You try using list comprehension with tuple unpacking here.

[val for idx, *rem in lst for val in rem[:idx]] 
# [1, 2, 8]

NB This solution assumes you would always have a sub-list of size 1 or greater. We can filter out empty sub-lists using filter(None, lst)

You can try List-comprehension:

>>> [sub[i] for sub in lst for i in range(1, sub[0]+1) ]
[1, 2, 8]

PS: The solution expects each sublist to be a non-empty list, else it will throw IndexError exception due to sub[0].

list1=[[2, 1, 2, 3, 4],
       [0, 4, 5],
       [1, 8, 9]]
list2= []

for nested_list in list1:
    for i in range(nested_list[0]):
        list2.append(nested_list[i+1])

Another list comprehension

sum([x[1:x[0] + 1] for x in arr], [])
# [1, 2, 8]

Using builtin function map to apply the picking function, and using itertools.chain to flatten the resulting list of list:

def pick(l):
    return l[1:1+l[0]]

ll = [[2, 1, 2, 3, 4], [0, 4, 5], [1, 8, 9]]

print( list(map(pick, ll)) )
# [[1, 2], [], [8]]

print( list(itertools.chain.from_iterable((map(pick, ll)))) )
# [1, 2, 8]

Or alternatively, with a list comprehension:

ll = [[2, 1, 2, 3, 4], [0, 4, 5], [1, 8, 9]]

print( [x for l in ll for x in l[1:1+l[0]]] )
# [1, 2, 8]

Two important notes:

  • I've renamed your list of lists ll rather than list. This is because list is already the name of the builtin class list in python. Shadowing the name of a builtin is very dangerous and can have unexpected consequences. I strongly advise you never to use the name of a builtin, when naming your own variables.

  • For both solutions above, the error-handling behaves the same: exception IndexError will be raised if one of the sublists is empty (because we need to access the first element to know how many elements to pick, so an error is raised if there is no first element). However, no exception will be raised if there are not enough elements in one of the sublists. For instance, if one of the sublists is [12, 3, 4], then both solutions above will silently pick the two elements 3 and 4, even though they were asked to pick 12 elements and not just 2. If you want an exception to be raised for this situation, you can modify function pick in the first solution:

def pick(l):
    if len(l) == 0 or len(l) <= l[0]:
        raise ValueError('in function pick: two few elements in sublist {}'.format(l))
    return l[1:1+l[0]]

ll = [[2, 1, 2, 3, 4], [0, 4, 5], [1, 8, 9], [12, 3, 4]]

print( [x for l in ll for x in l[1:1+l[0]]] )
# [1, 2, 8, 3, 4]

print( [x for l in ll for x in pick(l)] )
# ValueError: in function pick: two few elements in sublist [12, 3, 4]
Related