thomaso
thomaso

Reputation: 41

Recursive function to find the levels of elements in list

I need to write a recursive function that returns a list from original list where each element consists of elements in that level. For example [1, [2, [3], [[4], 3], 2], 1, []] for that list the result should be [[1, 1], [2, 2], [3, 3], [4]] (elements 1, 1 are on level 1, elements 2, 2 on level 2 and etc.). And another example would be [8, [6, 7, [-1], [4, [[10]]], 2], 1] turns to [[8, 1], [6, 7, 2], [-1, 4], [], [10]]. I've tried to go through every element and when it's not a list, add the element to result. Else I call out the function again until I reach just a element again. But the problem is that it just adds all the elements to result. I'm stuck and don't really know how to continue.

def levels(x, level = 0, result = []):
    for element in x:
        if type(element) is not list:
            result.append(element)
        else:
            levels(element, level + 1, result)

    return result

print(levels([1, [2, [3], [[4], 3], 2], 1, []]))
print(levels([8, [6, 7, [-1], [4, [[10]]], 2], 1]))

My result:

[1, 2, 3, 4, 3, 2, 1]
[1, 2, 3, 4, 3, 2, 1, 8, 6, 7, -1, 4, 10, 2, 1]

What it should be:

[[1, 1], [2, 2], [3, 3], [4]]
[[8, 1], [6, 7, 2], [-1, 4], [], [10]]

Any help is greatly appreciated.

Upvotes: 3

Views: 629

Answers (3)

skomlev
skomlev

Reputation: 27

This solution is very explicit and you can simplify and avoid some parameters.

It works by reference to sub-lists

def levels(my_list, index=0, new_list=None):
    if new_list is None:
        new_list = []

    if len(new_list) > index:
        items = new_list[index]
    else:
        items = []
        new_list.append(items)

    for item in my_list:
        if isinstance(item, list):
            nex_index = index + 1
            levels(item, nex_index, new_list)
        else:
            items.append(item)

    return new_list

Run:

> print(levels([8, [6, 7, [-1], [4, [[10]]], 2], 1]))
> [[8, 1], [6, 7, 2], [-1, 4], [], [10]]

> print(levels([1, 1, [], [2, 2, [3], [3, [4]]]]))
> [[1, 1], [2, 2], [3, 3], [4]]

Upvotes: 2

גלעד ברקן
גלעד ברקן

Reputation: 23955

Recursion usually means the type of value that's returned from later calls is the same as earlier calls, and the latter build on the former. So in this case, if we call the function on a new level, we'll get a list whose first element is representing one level deeper than the current.

def f(lst):
  result = [[]]
  deeper_res = []

  for e in lst:
    if isinstance(e, list):
      deeper = f(e)
      k = max(0, len(deeper) - len(deeper_res))
      deeper_res = [a + b for a,b in zip(deeper_res + [[]] * k, deeper)]
    else:
      result[0].append(e)

  return result + deeper_res

print(f([8, [6, 7, [-1], [4, [[10]]], 2], 1]))
# [[8, 1], [6, 7, 2], [-1, 4], [], [10]]
print(f([1, 1, [], [2, 2, [3], [3, [4]]]]))
#[[1, 1], [2, 2], [3, 3], [4]]

Upvotes: 1

Ajax1234
Ajax1234

Reputation: 71451

You can use recursion with itertools.groupby:

from itertools import groupby as gb
data = [1, [2, [3], [[4], 3], 2], 1, []] 
def get_levels(d, c=0):
  yield (c, [])
  for i in d:
     yield from ([(c, i)] if not isinstance(i, list) else get_levels(i, c+1))

r = [list(b) for _,b in gb(sorted(get_levels(data), key=lambda x:x[0]), key=lambda x:x[0])]
result = [list(filter(lambda x:x != [], [k for _, k in b])) for b in r]

Output:

[[1, 1], [2, 2], [3, 3], [4]]

Input test 2:

data = [8, [6, 7, [-1], [4, [[10]]], 2], 1]
...

Output:

[[8, 1], [6, 7, 2], [-1, 4], [], [10]]

Upvotes: 3

Related Questions