从深度嵌套列表中的数组中提取元素，保留列表结构

Question

我在问这个问题的一种概括：

Best way to extract elements from nested lists。

这也与这些问题有关：

recursive function for extract elements from deep nested lists/tuples

Scheme - find most deeply values nested lists

Scheme - find most deeply nested lists

基本上，我有一些任意的嵌套列表结构，在底部有各种形状相同的 numpy 数组。我想迭代或切片所有这些底层数组，同时保留它们所在的嵌套列表结构。这是关于在输出中保留嵌套结构的这一部分，这似乎在这些其他问题中没有得到回答。

所以，例如：

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

我们想要，示意性地，

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

或

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

我确信一些聪明的递归函数或其他东西可以做到这一点，我的大脑现在还没有想出它......

我想如果这将视图返回到底层数组而不是复制它们的一部分，那也是最好的。

编辑：也许这样的事情可以工作：https://stackoverflow.com/a/43357135/1447953。该方法需要以某种方式将 numpy 切片操作转换为函数，我想这可以根据具体情况进行，所以也许这就是要走的路。

Answer 1

也许是一个像这样的生成器：

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

如果第一个元素是 number 类型，将返回 ndarray。

Answer 2

快速的第一次尝试：

def recurse(f, alist):
    def foo(f, item):
        if isinstance(item,list):
            return recurse(f, item)
        else:
            return f(item)
    return [foo(f, item) for item in alist]

测试用例：

In [1]: A = np.arange(10)
In [2]: alist = [[A,A],[[A],[A,A]],A]

In [6]: recurse(lambda A: A[3], alist)
Out[6]: [[3, 3], [[3], [3, 3]], 3]

In [7]: recurse(lambda A: A[:3], alist)
Out[7]: 
[[array([0, 1, 2]), array([0, 1, 2])],
 [[array([0, 1, 2])], [array([0, 1, 2]), array([0, 1, 2])]],
 array([0, 1, 2])]

不限于索引：

In [10]: recurse(lambda A: A.sum(), alist)
Out[10]: [[45, 45], [[45], [45, 45]], 45]

Answer 3

感谢您的帮助！我想出了以下巫术来做我想做的事，部分基于https://stackoverflow.com/a/43357135/1447953。我还将它概括为允许具有镜像嵌套结构的二进制（和更高级别）操作。它基本上是“映射”的概括，适用于嵌套结构。可以更好地检查结构，例如它不能确保长度匹配，但我猜如果是这种情况，'map' 会抛出错误。

def apply_f(f,*iters):
    """Apply some function to matching 'bottom level' objects 
       in mirrored nested structure of lists,
       return the result in the same nested list structure.
       'iters' should be lists which have identical
       nested structure, and whose bottom level elements can
       be used as arguments to 'f'.
    """
    # We have to descend all list structures in lock-step!
    if all(isinstance(item,list) for item in iters):
        return list(map(lambda *items: apply_f(f,*items), *iters))
    elif any(isinstance(item,list) for item in iters):
        raise ValueError("Inconsistency in nested list structure of arguments detected! Nested structures must be identical in order to apply functions over them")
    else:
        return f(*iters)

可以这样使用：

A = np.ones((2,2))

a = [A,[A,A],A,[[A,A],A]]

def my_sum(a,b,c):
    return a + b + c

b = apply_f(my_sum,a,a,a)

print(a)
print(b)

# Check structure
print(apply_f(lambda A: A.shape, a))
print(apply_f(lambda A: A.shape, b))

输出：

[array([[1., 1.],
       [1., 1.]]), [array([[1., 1.],
       [1., 1.]]), array([[1., 1.],
       [1., 1.]])], array([[1., 1.],
       [1., 1.]]), [[array([[1., 1.],
       [1., 1.]]), array([[1., 1.],
       [1., 1.]])], array([[1., 1.],
       [1., 1.]])]]
[array([[3., 3.],
       [3., 3.]]), [array([[3., 3.],
       [3., 3.]]), array([[3., 3.],
       [3., 3.]])], array([[3., 3.],
       [3., 3.]]), [[array([[3., 3.],
       [3., 3.]]), array([[3., 3.],
       [3., 3.]])], array([[3., 3.],
       [3., 3.]])]]
[(2, 2), [(2, 2), (2, 2)], (2, 2), [[(2, 2), (2, 2)], (2, 2)]]
[(2, 2), [(2, 2), (2, 2)], (2, 2), [[(2, 2), (2, 2)], (2, 2)]]

当然，应用函数可以对底层数组做任何事情，无论是提取元素、切片、添加它们，等等。