To split a list into N chunks in Python, you can use iterators:
def split_list(lst, chunk_size):
for i in range(0, len(lst), chunk_size):
yield lst[i:i+chunk_size]
# Example use
lst = [1, 2, 3, 4, 5, 6, 7, 8, 9]
chunk_size = 3
for chunk in split_list(lst, chunk_size):
print(chunk)
Notice that it isn’t always possible to produce a list where the chunks are equal in length. This is because the list elements might not be evenly divisible to the N chunks.
This is a comprehensive guide to splitting a list into N chunks in Python. In case you’re looking for a quick solution, I’m sure the above will do. But if you’re learning Python, make sure to read the entire guide to figure out multiple approaches to splitting lists into chunks. The idea of this guide is not to use an existing solution but to implement the logic yourself.
Let’s jump into it!
1. Iterator Approach
Here is an example of how to split a Python list into equally-sized chunks using iterators and the yield keyword:
def split_list(lst, chunk_size):
for i in range(0, len(lst), chunk_size):
yield lst[i:i+chunk_size]
# Example usage
lst = [1, 2, 3, 4, 5, 6, 7, 8, 9]
chunk_size = 3
for chunk in split_list(lst, chunk_size):
print(chunk)
This code will output the following:
[1, 2, 3] [4, 5, 6] [7, 8, 9]
Let’s take a closer look at the code.
The split_list function takes two arguments: a list lst and an integer chunk_size that specifies the size of each chunk. The function uses a for loop to iterate over the list, and yields a sublist of lst starting at the current index and ending at the current index plus the chunk size.
Here’s a step-by-step explanation of how the code works:
- The
split_listfunction is called with a listlstand a chunk sizechunk_size. - The
forloop iterates over the list, starting at index 0 and incrementing the index bychunk_sizeeach time. For example, if the chunk size is 3, the first iteration will start at index 0, the second iteration will start at index 3, the third at 6, and so on. - On each iteration, the
yieldkeyword constructs a sublist oflststarting at the current index and ending at the current index plus the chunk size. E.g. if the current index is 0 and the chunk size is 3, the sublist will belst[0:3], which is the three elements oflst. - The
forloop in the example code then iterates over the chunks yielded bysplit_listand prints each chunk.
The result of all of this is a list that is split into N chunks.
But What On Earth Does the ‘yield’ Do?
The previously introduced approach is the easiest one to split a list into chunks. But if you’re new to the yield keyword and iterators, this solution might just leave you confused.
In case you’re interested in iterators, make sure to read this complete guide to iterators and iterables in Python.
The next section teaches you how to do the previous approach without using the yield keyword but using lists instead.
2. For Loop Approach
Here’s another way to split lists into chunks in Python. This approach uses for loops and lists and is thus a bit more beginner-friendly than the previous example.
Here are the steps you need to take:
- Determine the number of chunks you want to split the list into. Let’s call this number
n. - Floor-divide the length of the list by
nto find the size of each chunk (floor division rounds down so that the chunk size isn’t e.g. 3.3333 but instead just 3). We’ll call this numberchunk_size. - Use the
range()function to create a list of numbers that specify the indexes where each chunk should start. If the original list has 10 elements and you want to split it into 3 chunks, the list of starting indexes would be[0, 3, 6]. - Use a
forloop to iterate over the list of starting indexes, and use thelist[start:end]syntax to extract each chunk from the original list.
Here’s the code that implements those steps. Make sure to read the comments to stay on top of it all the time!
# Set the number of chunks N = 4 # Create a list of numbers my_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] # Create an empty list to store the chunks chunks = [] # Iterate over the elements of the list in groups of N for i in range(0, len(my_list), N): # Extract each group of N elements as a sublist chunk = my_list[i:i + N] # Append the sublist to the list of chunks chunks.append(chunk) # Print the chunks print(chunks)
This code produces the following output:
[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10]]
2.1. List Comprehensions Can Do the Same
When you’re dealing with for loops, in some cases, you might be able to use a list comprehension to tidy up the code.
It’s always up for debate as to whether you should use comprehension or not. This is because a list comprehension for sure shortens the code but might actually make it less readable.
Anyway, here’s the previous example that uses a list comprehension instead of a for loop:
# Define the input list input_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] # Define the chunk size chunk_size = 3 # Create the output list using a list comprehension output_list = [input_list[i:i + chunk_size] for i in range(0, len(input_list), chunk_size)] # Print the output list print(output_list)
This will produce the following output:
[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]
3. While Loop
In the previous examples, you used a for loop to split a list into chunks. Because you can use a for loop, you can certainly do one with a while loop too!
Here’s what the code looks like when using a while loop:
# define the list of elements
my_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
# define the size of the chunks
chunk_size = 3
# create an empty list to store the chunks
chunks = []
# create an index variable to track the current position in the list
index = 0
# loop until the index is larger than the length of the list
while index < len(my_list):
# get the sublist of the current chunk
chunk = my_list[index:index+chunk_size]
# append the chunk to the list of chunks
chunks.append(chunk)
# update the index to the next chunk
index += chunk_size
# print the resulting chunks
print(chunks)
Output:
[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]
This solution creates an empty list to store the resulting chunks and then uses a while loop to iterate over the original list by chunk sizes, appending each chunk to the list of chunks. The index variable is used to track the current position in the list and is updated at each iteration to the next chunk.
4. Use NumPy to Split into N Chunks
In Python, there’s a popular math and science module called NumPy that the scientific Python community uses a lot.
If you’re using NumPy already, you can use the array_split() function to split lists into chunks. This function takes in the list as the first argument and the size of the chunks as the second argument.
Notice that this approach forces the elements into N chunks. The leftover values aren’t placed in their own chunk but are pushed to the last chunk instead.
For example, let’s split a list of 10 numbers into 3 chunks:
import numpy as np # Define the list to be split my_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] # Split the list into chunks of size 3 chunks = np.array_split(my_list, 3) # Print the resulting chunks print(chunks) # Output: [array([1, 2, 3]), array([4, 5, 6]), array([7, 8, 9, 10])]
Notice how the last chunk now has 4 elements instead of taking the extra element and placing it in its own chunk.
Thanks for reading. Happy coding!