展平任意深度的嵌套列表的空间和时间复杂度

Question

给定一个包含任意嵌套级别的嵌套列表的Python列表，目标是返回一个完全扁平化的列表，即对于样本输入

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

，输出应该是

[1, 2, 3, 1]

。

我的解决方案：

def flatten(lst):
  stack = [[lst, 0]]
  result = []

  while stack:
    current_lst, start_index = stack[-1]
    for i in range(start_index, len(current_lst)):
      if isinstance(current_lst[i], list):
        # Update the start_index of current list
        # to the next element after the nested list
        stack[-1][1] = i + 1
        # Add nested list to stack
        # and initialize its start_index to 0
        stack.append([current_lst[i], 0])
        break
      # non list item
      # add item to result
      result.append(current_lst[i])
    else:
      # no nested list
      # remove current list from stack
      stack.pop()

  return result

我想知道我的解决方案的时间和空间复杂度（辅助空间）是否正确。

我的想法

时间复杂度：

我相信该解决方案的时间复杂度为O(m + n)，其中m是各级嵌套列表的总数，n是原子元素的总数所有级别（非列表元素）

空间复杂度：

我相信空间复杂度是O(d)，其中d是最嵌套列表的深度。这是因为堆栈跟踪当前的遍历状态，其大小与嵌套深度成正比

解决方案正确吗？

时空分析正确吗？

Answer 1

是的，解决方案是正确的。

是的，时空分析是正确的……如果不把

result

使用的空间算作辅助空间的话，也是合理的。尽管请注意

result

会过度分配/重新分配，您可以将其视为占用 O(n) 辅助空间。您可以通过对整个输入执行两次传递来优化它，一次计算原子元素，然后创建

result = [None] * n

，然后另一次传递来填充它。

顺便说一句，最好使用迭代器而不是你自己的列表+索引对：

def flatten(lst):
  stack = [iter(lst)]
  result = []
  while stack:
    for item in stack[-1]:
      if isinstance(item, list):
        stack.append(iter(item))
        break
      result.append(item)
    else:
      stack.pop()
  return result

在线尝试！

或者使用提到的空间优化：

def flatten(lst):
  def atoms():
    stack = [iter(lst)]
    while stack:
      for item in stack[-1]:
        if isinstance(item, list):
          stack.append(iter(item))
          break
        yield item
      else:
        stack.pop()
  n = sum(1 for _ in atoms())
  result = [None] * n
  for i, result[i] in enumerate(atoms()):
    pass
  return result

print(flatten([1, [2], [[[3]]], 1]))

在线尝试！