How to get all elements at a specific depth-level in a nested list?

Question

I am looking for a method to get all the elements nested at a user-defined list depth level eg:

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

# example 1
level = 1  # user defined level
output = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]

# example 2
level = 2
output = [[1, 2], [3, 4], [5, 6], [7, 8]]

# example 3
level = 3
output = [1, 2, 3, 4, 5, 6, 7, 8]

Answer 1

You can just use a recursive algorithm, for example:

output = []
def extract(lists, d):
    if d == 1:
        return output.extend(lists)

    for sub_list in lists:
        extract(sub_list, d - 1)

For level 1:

extract(lst, 1)
print(output)
>>> [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]

For level 2:

extract(lst, 2)
print(output)
>>> [[1, 2], [3, 4], [5, 6], [7, 8]]

For level 3

extract(lst, 3)
print(output)
>>> [1, 2, 3, 4, 5, 6, 7, 8]

Answer 2

You can use chain.from_iterable to go down one level every time:

from itertools import chain

def get_at_level(lst, level):
    for _ in range(level-1):
        lst = chain.from_iterable(lst)

    return list(lst)

Examples:

>>> lst = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
>>> get_at_level(lst, 1)
[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
>>> get_at_level(lst, 2)
[[1, 2], [3, 4], [5, 6], [7, 8]]
>>> get_at_level(lst, 3)
[1, 2, 3, 4, 5, 6, 7, 8]

Please NOTE , that the function returns only a shallow copy of the original list. So assuming you call with any level that is not the lowest - you will have the same references to sub-lists from the original. This means that modifying the returned list might * change the original, If you don't care about the original list. that's fine, If you don't want it to be changed, create a deep copy of it in the first line of the function.

---

^{* Changing the first level of the returned will not be a problem because as explained, list returns a shallow copy. BUT, doing something like get_at_level(lst, 2)[0][0] = 0 will also affect the original.}

Answer 3

This isn't a very clean solution for the problem, but it works!

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

def function(lst, level):
    if level == 1:
        print(lst)

    elif level == 2:
      l2 = []
      for x in lst:
        for i in x:
            l2.append(i)
      print(l2)

    elif level == 3:
      l3 = []
      for x in lst:
        for i in x:
            for a in i:
                l3.append(a)
      print(l3)

    else:
      print("Invalid depth level")

function(lst, 1) #level 1, 2 or 3

The problem here is that it isn't dynamic.

Answer 4

A recursive function like func as below would work. It is basically the same as marcos' answer.

func accepts a list as its first argument and depth as the second.

from functools import reduce

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


func = lambda x, d: x if d == 1 else func(reduce(lambda a,b: a+b, x), d-1)

func(lst, 3) # output [1,2,3,4,5,6,7,8]

Answer 5

You can use a recursive generator function:

def flatten(d, l=1):
   for i in d:
      yield from ([i] if l == 1 else flatten(i, l-1))

lst = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
print(list(flatten(lst, 3)))
print(list(flatten(lst, 2)))
print(list(flatten(lst, 1)))

Output:

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