5 Best Ways to Convert a 1D List to a 2D List of Variable Length in Python

Rate this post

πŸ’‘ Problem Formulation: Converting a one-dimensional list into a two-dimensional list of variable lengths in Python is a common task that developers encounter. This involves transforming a flat list into a nested list, where each sub-list may contain a different number of elements. For instance, the input [1, 2, 3, 4, 5, 6] could be converted into [[1, 2], [3, 4], [5, 6]] or [[1], [2, 3, 4], [5, 6]] depending on the partitioning logic.

Method 1: Using Loops to Manually Split the List

This method involves manually iterating over the elements of the list and creating a new sublist whenever the desired condition is met. It is very flexible as it allows you to implement complex logic to determine sublist lengths.

Here’s an example:

def manual_split(lst, sizes):
    result = []
    it = iter(lst)
    for size in sizes:
        result.append([next(it) for _ in range(size)])
    return result

original_list = [1, 2, 3, 4, 5, 6]
sizes = [2, 1, 3]
new_list = manual_split(original_list, sizes)
print(new_list)

Output:

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

We define a function manual_split() that accepts a list and a list of sizes. It iterates through the original list, creating sublists of the lengths specified in the sizes list. This method is basic but offers great control over the resulting sublists’ shapes.

Method 2: Using List Comprehensions with Slicing

List comprehensions provide a concise way to create lists. Combined with slicing, this method allows for creating sublists of fixed lengths from a 1D list, with the last sublist possibly being shorter if there aren’t enough elements left.

Here’s an example:

original_list = [1, 2, 3, 4, 5, 6]
size = 2
new_list = [original_list[i:i + size] for i in range(0, len(original_list), size)]
print(new_list)

Output:

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

This snippet uses a list comprehension to create a new 2D list where each sublist has a fixed size, defined by size. The range() function is used to step through the original list by the size of each chunk.

Method 3: Using the numpy.array_split() Method

The NumPy library provides a convenient method array_split() to split an array into multiple sub-arrays of equal or near-equal size. This method is optimal for scientific computing tasks where NumPy arrays are commonly used.

Here’s an example:

import numpy as np

original_list = [1, 2, 3, 4, 5, 6]
num_of_sublists = 3
new_list = np.array_split(original_list, num_of_sublists)
print(new_list)

Output:

[array([1, 2]), array([3, 4]), array([5, 6])]

This code uses NumPy’s array_split() function to split the original list into num_of_sublists sublists. NumPy handles the distribution of elements, attempting to balance the sublist sizes as evenly as possible.

Method 4: Using itertools

The itertools library has a tool called grouper() (it has to be implemented) that groups elements from an iterable into fixed-length chunks, filling missing values with None if the chunk is incomplete. This method is ideal for handling large datasets efficiently.

Here’s an example:

from itertools import zip_longest

def grouper(iterable, n, fillvalue=None):
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

original_list = [1, 2, 3, 4, 5, 6, 7]
size = 3
new_list = list(grouper(original_list, size))
print(new_list)

Output:

[(1, 2, 3), (4, 5, 6), (7, None, None)]

In this snippet, the grouper() function from itertools is used to divide the list into sublists of equal length. Incomplete sublists are filled with None to maintain a consistent size across sublists.

Bonus One-Liner Method 5: The Functional Approach with map() and iter()

This one-liner employs Python’s map function alongside an iterator to partition the list into variable-length sublists defined in a separate sizes list, making it a compact and elegant solution.

Here’s an example:

original_list = [1, 2, 3, 4, 5, 6]
sizes = [2, 3, 1]
it = iter(original_list)
new_list = list(map(lambda size: list(islice(it, size)), sizes))
print(new_list)

Output:

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

This one-liner uses the map() function to apply an inline lambda that slices the list into sublists of sizes specified in sizes list. The iter() ensures that we continue slicing from the last position in the original list.

Summary/Discussion

  • Method 1: Manual Splitting with Loops. Offers the most control for complex scenarios. Can be cumbersome for large datasets or performance-critical applications.
  • Method 2: List Comprehensions with Slicing. Clean and Pythonic. Limited to fixed-size chunks. Inefficient if the size of sublists varies significantly.
  • Method 3: NumPy’s array_split(). Easy for equal or nearly equal split. Requires NumPy dependency, which might be unnecessary for simple tasks.
  • Method 4: Using itertools. Efficient and elegant, especially for large data. The presence of None for padding may require additional cleanup.
  • Bonus Method 5: Functional Approach with map() and iter(). Compact one-liner. Might be less readable for beginners. The sizes list must be accurately predefined.