Python – Extracts elements from arrays in deeply nested lists, preserving list structure

I’m asking a generalization of this question:

Best way to extract elements from nested lists .

It’s also somehow related to these issues:

recursive function for extract elements from deep nested lists/tuples

Scheme – find most deeply values nested lists

Scheme – find most deeply nested lists

Essentially, I have some arbitrary nested list structure with various identically shaped numpy arrays at the bottom. I want to iterate over or slice all these underlying arrays while preserving the nested list structure they are in. This part is about preserving nested structures in the output, and these other questions don’t seem to be answered.

So, for example:

A = np.ones((3,3,3))
nest = [[A,A,A],[[A,A,[A,A]],[A,A]],A]

We want, schematically,

nest[0,...] == [[A[0,...],A[0,...],A[0,...]],[[A[0,...],A[0,...],[A[0,...],A[0,...]]],[A[0,...],A[0,...]]],A[0,...]]

or

nest[1:3,5,:] == [[A[1:3,5,:],A[1:3,5,:],A[1:3,5,:]],[[A[1:3,5,:],A[1:3,5,:],[A[1:3,5,:],A[1:3,5,:]]],[A[1:3,5,:],A[1:3,5,:]]],A[1:3,5,:]]

I’m sure some clever recursive function or something else could do that my brain hasn’t figured out right now….

I guess it would be best if this returned the View to the underlying array instead of copying part of it.

EDIT: Maybe something like this would work: https://stackoverflow.com/a/43357135/1447953 The method needs to somehow convert the numpy slice operation to a function, which I guess can be done on a case-by-case basis, so maybe that’s the way to go.

Maybe a generator like this :

def get_nested_elements(elements):
    if not elements or isinstance(elements[0], np.number):
        yield elements
    else:
        for node in elements:
            for e in get_nested_elements(node):
                yield e

If the first element is a numeric type, an ndarray is returned.