How to Flatten Nested Lists of Arbitrary Depth in Python

Dealing with nested lists in Python can become tricky, especially when you need to process all elements regardless of how deeply nested they are.
This article explores different methods to flatten a nested list in Python, regardless of its depth. We’ll cover iterative and recursive approaches, ensuring you understand the core concepts behind each. We will also use list comprehension to flatten the lists.

Let’s consider a sample nested list:

nested_list = [1, [2, [3, 4], 5], 6]
print(nested_list)

[1, [2, [3, 4], 5], 6]

Method 1: Recursive Flattening

This approach uses recursion to traverse the nested list. If an element is a list, the function calls itself to flatten that sublist. Otherwise, it appends the element to the result.

def flatten_recursive(nested_list):
    flat_list = []
    for element in nested_list:
        if isinstance(element, list):
            flat_list.extend(flatten_recursive(element))
        else:
            flat_list.append(element)
    return flat_list

nested_list = [1, [2, [3, 4], 5], 6]
flat_list = flatten_recursive(nested_list)
print(flat_list)

[1, 2, 3, 4, 5, 6]

Explanation: The flatten_recursive function checks each element. If an element is a list, it recursively calls itself to flatten that sublist. The extend method adds all elements from the flattened sublist to the flat_list. If the element is not a list, it’s simply appended to flat_list.

Method 2: Iterative Flattening with a Stack

This method uses a stack to keep track of the elements to be processed. It avoids recursion and can be more memory-efficient for very deep nested lists.

def flatten_iterative(nested_list):
    flat_list = []
    stack = [nested_list]
    while stack:
        current = stack.pop()
        if isinstance(current, list):
            stack.extend(current[::-1])  # Reverse to maintain order
        else:
            flat_list.append(current)
    return flat_list

nested_list = [1, [2, [3, 4], 5], 6]
flat_list = flatten_iterative(nested_list)
print(flat_list)

[1, 2, 3, 4, 5, 6]

Explanation: The flatten_iterative function uses a stack to keep track of elements to process. When a list is encountered, its elements are added to the stack in reverse order (current[::-1]) to ensure the original order is preserved during flattening. Non-list elements are appended to the flat_list.

Method 3: Using List Comprehension (for shallow nesting)

For lists with a known, limited depth, list comprehension offers a concise way to flatten. This example assumes a maximum nesting depth of 2.

nested_list = [1, [2, 3], [4, [5, 6]]] #limited depth

flat_list = [item for sublist in nested_list for item in (sublist if isinstance(sublist, list) else [sublist])]
print(flat_list)

#Flatten one level deeper (if needed for above example, not recommended in general case).
flat_list_deep = [item2 for sublist in nested_list for item in (sublist if isinstance(sublist, list) else [sublist]) for item2 in (item if isinstance(item, list) else [item])]
print(flat_list_deep)

[1, 2, 3, 4, [5, 6]]
[1, 2, 3, 4, 5, 6]

Explanation: The outer loop iterates through the main list. The inner loop iterates through the sublist (or a single-element list containing the item if it’s not a list). This approach is not recommended for arbitrary depth because extending it becomes unreadable and complex. The second example shows how quickly this can become unwieldy.

Method 4: Using itertools.chain.from_iterable

This method leverages the `itertools` library for a more functional approach to flattening. It avoids recursion and can be quite efficient.

import itertools

def flatten_itertools(nested_list):
    def gen(lst): # Generator function
      for item in lst:
          if isinstance(item, list):
              yield from gen(item) # recursive call using yield from
          else:
              yield item

    return list(gen(nested_list)) #Convert generator to a list

nested_list = [1, [2, [3, 4], 5], 6]
flat_list = flatten_itertools(nested_list)
print(flat_list)

[1, 2, 3, 4, 5, 6]

Explanation: The code defines a generator function `gen` that recursively yields elements from the nested list. The `yield from` statement is used to yield all values from the recursive calls of the generator for sublists. The main function converts the generator to a list.

Method 5: Using a while loop to check if list has other lists inside.

This method uses a while loop to flatten the list, by looping until no lists are found.

def flatten_while_loop(nested_list):
  while True:
    flat_list = []
    list_changed = False
    for element in nested_list:
      if isinstance(element, list):
        flat_list.extend(element)
        list_changed = True
      else:
        flat_list.append(element)
    nested_list = flat_list
    if not list_changed:
      break
  return nested_list

nested_list = [1, [2, [3, 4], 5], 6]
flat_list = flatten_while_loop(nested_list)
print(flat_list)

[1, 2, 3, 4, 5, 6]

Explanation: The function uses a while loop, which continues while list_changed is True. Inside the loop, the code iterates through each element in the nested list. If the element is a list, its content is extended into the flat_list and list_changed is set to True. Otherwise, the element is directly appended to the flat_list. After the loop completes (when no lists are found), the nested_list is updated with the content of flat_list and returned.