How to get all array elements at depth in python?-CodePudding

I am trying to make a function which takes an array and returns the array sorted by the depth:

sortArrDepth([1, [5, 6], [4, 67, [34]], 7])

Would return

[[1, 7], [5, 6, 4, 67], [34]]

However, it has to be able handle arrays with any max depth and I would prefer to not use any external modules. If it helps, here is a function to get the max depth:

def getAstDepth(ast):
    depth = 0
    for i in last:
        if isinstance(i, list):
            depth = max(getAstDepth(i), depth)
    return depth

Any help appreciated!

CodePudding user response：

Here is one way using a recursive function:

def sortArrDepth(l, out=[], depth=0):
    if len(out)<=depth:
        out  = [[] for _ in range(depth-len(out) 1)]
    for v in l:
        if isinstance(v, list):
            sortArrDepth(v, out=out, depth=depth 1)
        else:
            out[depth].append(v)
    return out

example:

>>> sortArrDepth([1, [5, 6], [4, 67, [34]], 7])
[[1, 7], [5, 6, 4, 67], [34]]

>>> sortArrDepth([1, [5, 6], [[[0]]], [4, 67, [34]], 7])
[[1, 7], [5, 6, 4, 67], [34], [0]]

CodePudding user response：

You can use itertools.zip_longest and recursion to get at the deeper levels:

from itertools import zip_longest
def sort_by_depth(seq):
    first = [item for item in seq if not isinstance(item, list)]
    nested = [sort_by_depth(item) for item in seq if isinstance(item, list)]
    return [first]   [sum(item, start=[]) for item in zip_longest(*nested, fillvalue=[])]