TL;DR
此实现处理具有无限深度的嵌套迭代:
def enumerate_items_from(iterable):
cursor_stack = [iter(iterable)]
item_index = -1
while cursor_stack:
sub_iterable = cursor_stack[-1]
try:
item = next(sub_iterable)
except StopIteration:
cursor_stack.pop()
continue
if len(cursor_stack) == 1:
item_index += 1
if not isinstance(item, str):
try:
cursor_stack.append(iter(item))
continue
except TypeError:
pass
yield item, item_index
def flat(iterable):
return map(list, zip(*enumerate_items_from(a)))
可用于产生所需的输出:
>>> nested = ['', '', '', ['', [['a', 'b'], ['c']]], [[['a', 'b'], ['c']]], [[['d']]]]
>>> flat_list, item_indexes = flat(nested)
>>> print(item_indexes)
[0, 1, 2, 3, 3, 3, 3, 4, 4, 4, 5]
>>> print(flat_list)
['', '', '', '', 'a', 'b', 'c', 'a', 'b', 'c', 'd']
请注意,您可能应该将索引放在首位,以模仿 enumerate 所做的事情。对于已经知道enumerate的人来说会更容易使用。
重要说明 除非您确定您的列表不会嵌套太多,否则不应使用任何基于递归的解决方案。否则,一旦您的嵌套列表深度大于 1000,您的代码就会崩溃。我解释这个here。请注意,对str(list) 的简单调用将在depth > 1000 的测试用例上崩溃(对于某些python 实现,它不止于此,但它始终是有界的)。使用基于递归的解决方案时,您将遇到的典型异常是(简而言之,这是由于 python 调用堆栈的工作方式):
RecursionError: maximum recursion depth exceeded ...
实现细节
我一步一步来,首先我们将一个列表展平,然后将展平后的列表和所有项目的深度都输出,最后在“main”中输出列表和对应的项目索引列表”。
扁平化列表
话虽如此,这实际上很有趣,因为迭代解决方案是为此完美设计的,您可以采用简单的(非递归)列表展平算法:
def flatten(iterable):
return list(items_from(iterable))
def items_from(iterable):
cursor_stack = [iter(iterable)]
while cursor_stack:
sub_iterable = cursor_stack[-1]
try:
item = next(sub_iterable)
except StopIteration: # post-order
cursor_stack.pop()
continue
if isinstance(item, list): # pre-order
cursor_stack.append(iter(item))
else:
yield item # in-order
计算深度
我们可以通过查看堆栈大小来访问深度,depth = len(cursor_stack) - 1
else:
yield item, len(cursor_stack) - 1 # in-order
这将返回对 (item, depth) 的迭代,如果我们需要将此结果分成两个迭代器,我们可以使用 zip 函数:
>>> a = [1, 2, 3, [4 , [[5, 6], [7]]], [[[8, 9], [10]]], [[[11]]]]
>>> flatten(a)
[(1, 0), (2, 0), (3, 0), (4, 1), (5, 3), (6, 3), (7, 3), (8, 3), (9, 3), (10, 3), (11, 3)]
>>> flat_list, depths = zip(*flatten(a))
>>> print(flat_list)
(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)
>>> print(depths)
(0, 0, 0, 1, 3, 3, 3, 3, 3, 3, 3)
我们现在将做一些类似的事情来使用项目索引而不是深度。
计算项目索引
要改为计算项目索引(在主列表中),您需要计算到目前为止所看到的项目数,这可以通过每次迭代一个 item_index 来完成深度为 0 的项目(当堆栈大小等于 1 时):
def flatten(iterable):
return list(items_from(iterable))
def items_from(iterable):
cursor_stack = [iter(iterable)]
item_index = -1
while cursor_stack:
sub_iterable = cursor_stack[-1]
try:
item = next(sub_iterable)
except StopIteration: # post-order
cursor_stack.pop()
continue
if len(cursor_stack) == 1: # If current item is in "main" list
item_index += 1
if isinstance(item, list): # pre-order
cursor_stack.append(iter(item))
else:
yield item, item_index # in-order
类似地,我们将使用 ˋzip, we will also use ˋmap 将两个迭代器拆分为两个迭代器中的对,以将两个迭代器转换为列表:
>>> a = [1, 2, 3, [4 , [[5, 6], [7]]], [[[8, 9], [10]]], [[[11]]]]
>>> flat_list, item_indexes = map(list, zip(*flatten(a)))
>>> print(flat_list)
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
>>> print(item_indexes)
[0, 1, 2, 3, 3, 3, 3, 4, 4, 4, 5]
改进——处理可迭代的输入
能够采用更广泛的嵌套迭代作为输入可能是可取的(特别是如果您构建它以供其他人使用)。例如,如果我们将嵌套的可迭代对象作为输入,则当前实现无法按预期工作,例如:
>>> a = iter([1, '2', 3, iter([4, [[5, 6], [7]]])])
>>> flat_list, item_indexes = map(list, zip(*flatten(a)))
>>> print(flat_list)
[1, '2', 3, <list_iterator object at 0x100f6a390>]
>>> print(item_indexes)
[0, 1, 2, 3]
如果我们想让它起作用,我们需要小心一点,因为字符串是可迭代的,但我们希望它们被视为原子项目(而不是作为字符列表)。而不是像我们之前那样假设输入是一个列表:
if isinstance(item, list): # pre-order
cursor_stack.append(iter(item))
else:
yield item, item_index # in-order
我们不会检查输入类型,而是尝试使用它,就像它是可迭代的一样,如果失败,我们将知道它不是可迭代的(鸭子类型):
if not isinstance(item, str):
try:
cursor_stack.append(iter(item))
continue
# item is not an iterable object:
except TypeError:
pass
yield item, item_index
通过这个实现,我们有:
>>> a = iter([1, 2, 3, iter([4, [[5, 6], [7]]])])
>>> flat_list, item_indexes = map(list, zip(*flatten(a)))
>>> print(flat_list)
[1, 2, 3, 4, 5, 6, 7]
>>> print(item_indexes)
[0, 1, 2, 3, 3, 3, 3]
构建测试用例
如果需要生成深度较大的测试用例,可以使用这段代码:
def build_deep_list(depth):
"""Returns a list of the form $l_{depth} = [depth-1, l_{depth-1}]$
with $depth > 1$ and $l_0 = [0]$.
"""
sub_list = [0]
for d in range(1, depth):
sub_list = [d, sub_list]
return sub_list
您可以使用它来确保我的实现在深度较大时不会崩溃:
a = build_deep_list(1200)
flat_list, item_indexes = map(list, zip(*flatten(a)))
我们也可以通过str函数检查我们不能打印这样的列表:
>>> a = build_deep_list(1200)
>>> str(a)
RecursionError: maximum recursion depth exceeded while getting the repr of an object
函数repr 由str(list) 调用来自输入list 的每个元素。
结束语
最后我同意递归实现更容易阅读(因为调用堆栈为我们做了一半的工作),但是在实现这样的低级函数时,我认为拥有这样的代码是一项很好的投资适用于所有情况(或至少您能想到的所有情况)。特别是当解决方案不是那么难的时候。这也是一种不要忘记如何编写在树状结构上工作的非递归代码的方法(除非您自己实现数据结构,否则这种情况可能不会经常发生,但这是一个很好的练习)。
请注意,我所说的“反对”递归的一切都是正确的,因为 python 在面临递归时不会优化调用堆栈的使用:Tail Recursion Elimination in Python。而许多编译语言都使用Tail Call recursion Optimization (TCO)。这意味着即使你在 python 中编写了完美的tail-recursive 函数,它也会在深度嵌套的列表上崩溃。
如果您需要有关列表展平算法的更多详细信息,可以参考我链接的帖子。