Sometimes we need to split the list of items into batches of specified size. In this tutorial, we’ll discuss a method to split list into batches using a generator in Python.
Let’s consider we have a list of numbers. We will split this list into a small batch of size 3.
a = [1,2,3,4,5,6,7,8,10]
So the output would be like this:
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
[10]
Batch generator implementation
It is easy to create a generator function in Python. Generator functions uses yield statement instead of return statement. A simple generator function as shown below.
def gen(): # defines a generator function
yield 42
The batch generator implementation has the following steps:
- Get the list and batch size.
- Get the start index for the batch using the range function.
- Slice the list using the start index and end index.
- Yield the result.
The generator function for the batch function can be defined as:
def generate_batch(lst, batch_size):
""" Yields batch of specified size """
for i in range(0, len(lst), batch_size):
yield lst[i : i + batch_size]
Let’s examine this generator function line by line.
for i in range(0, len(lst), batch_size):
Here range function used to get the start index for the list slice.
yield lst[i : i + batch_size]
Here list is sliced using the start index (i) and end index using ( i + batch_size).
test_list = list(range(1,11))
gen = generate_batch(test_list,3)
for batch in gen:
print(batch)
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
[10]
For this function, batch_size must be greater than 0. Improved function as shown.
def generate_batch(lst, batch_size):
""" Yields bacth of specified size """
if batch_size<=0:
return
for i in range(0, len(lst), batch_size):
yield lst[i : i + batch_size]
Conclusion
There you have it, a simple way to split the list into batches using a generator in Python. Read about list append extend performance.