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 e.g.:

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]

What did you try so far? Did you have a look at recursive functions? Implementing a recursive function extracting the sublists up to a given level might solve your issue. — albert, Apr 04 '21 at 22:11

score 2 · Answer 1 · answered Apr 04 '21 at 22:47

2

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]

answered Apr 04 '21 at 22:47

marcos

4,473
1
10
24

1

You need to mention that it also requires to "reset" `output` before every call... It is better to have it defined ***inside*** the function – Tomerikoo Apr 04 '21 at 23:00

Tomerikoo · Answer 2 · 2021-04-05T08:54:56.930

1

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.}

edited Apr 05 '21 at 08:54

answered Apr 04 '21 at 22:51

Tomerikoo

18,379
16
47
61

Why do we have to have the `ref` variable? From what I can see it isn't mutated only re-assigned, which wouldn't affect the caller's list. Please correct me if I'm wrong. – Roy Cohen Apr 04 '21 at 23:11
@RoyCohen I was just wondering the same thing... Check my edit :) – Tomerikoo Apr 04 '21 at 23:12
1

I tested this with the `list` call in the end of the function instead of inside the loop and it work. Please accept my edit. – Roy Cohen Apr 04 '21 at 23:16
Why do we need the `deepcopy`? I tested the code without it and it work fine. – Roy Cohen Apr 04 '21 at 23:21
@RoyCohen The function returns a list, possibly nested. The `list` constructor only creates a shallow copy and so if someone will change the returned list it might change the original and that might not be wanted. – Tomerikoo Apr 04 '21 at 23:22
@RoyCohen By the way thanks for your suggestion, I accidentally removed your edit because I was editing at the same time – Tomerikoo Apr 04 '21 at 23:23
1

Regarding the mutability problem, from what I see *every* other answer has this issue. Perhaps there should be a note saying something like "the `deepcopy` is only needed to stop mutation in the output effect the input. If that's not a concern, you may remove it." After all, a deep copy can be a potentially relativly heavy operation, especialy if `level` is small and `lst` contains a lot of levels. – Roy Cohen Apr 04 '21 at 23:44
@RoyCohen Thank you! That's a good point! I will edit that in – Tomerikoo Apr 05 '21 at 08:14

pupspulver · Answer 3 · 2021-04-04T22:54:40.073

0

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.

edited Apr 04 '21 at 22:54

answered Apr 04 '21 at 22:36

pupspulver

124
1
12

1

What about 4 levels? 5? 10? – Tomerikoo Apr 04 '21 at 22:42
Yeah that's why it isn't a very clean way of solving the problem, because is isn't dynamic. And in this case u don't have more than 3 levels. – pupspulver Apr 04 '21 at 22:50

score 0 · Answer 4 · answered Apr 04 '21 at 22:51

0

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]

answered Apr 04 '21 at 22:51

Shihao Xu

721
1
7
14

You shouldn't assign to a lambda function, see my suggested edit. – Roy Cohen Apr 04 '21 at 23:18

Ajax1234 · Answer 5 · 2021-04-04T23:36:31.500

0

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

edited Apr 04 '21 at 23:36

answered Apr 04 '21 at 23:16

Ajax1234

69,937
8
61
102

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

5 Answers5