The Art of Writing Loops in Python

Simple is better than complex

Yang Zhou
Yang Zhou
May 3 · 4 min read
Photo by Jie on Unsplash

The for loop is a very basic control flow tool of most programming languages. For example, a simple for loop in C looks like the following:

There are no other ways to write a for loop more elegantly in C. For a complex scenario, we usually need to write ugly nested loops or define lots of assistant variables (like the i in the above code).

Fortunately, when it comes to Python, things become much more convenient. We have many tricks to write much more elegant loops and they do make our lives easier. In Python, nested loops are not inevitable, assistant variables are not necessary, and we can even customise a for loop by ourselves.

This article will introduce some of the most helpful tricks for writing loops in Python. Hopefully, it can help you feel the beauty of Python.

Get Indexes and Values at Once

A common scenario of using a for loop is to get indexes and values from a list. When I started learning Python, I wrote my code like the following:

It works of course. But not Pythonic enough. After a few months, I got the standard Pythonic way to do this:

As shown above, the built-in enumerate function can make our lives easier.

Avoid Nested Loops by the Product Function

Nested loops are headaches. They can reduce the readability of our code and make things complex. For example, breaking out of the nested loops is usually not very easy. We need to know when the inner-most loop was broken, when the second inner-most loop was broken, and so on.

Fortunately, there is an awesome function called product from the built-in itertools module in Python. We can use it to avoid writing lots of nested loops.

Let’s feel how useful it is by a simple example:

As shown above, we need three nested loops to get three numbers whose sum is equal to 2077 from three lists. The code is not neat at all.

How about using the product function?

As shown above, with the help of the product function, only one loop is needed.

Because the product function generates the Cartesian product of input iterables. It can help us avoid nested loops in lots of scenarios.

Use the Itertools Module To Write Fancy Loops

In fact, the product function is just the tip of the iceberg. If you explore the built-in itertools module in Python. A new world will be open to you. This toolbox contains many useful methods to satisfy our needs about loops. The full list of them can be found on the official document. Let’s enjoy a few interesting usages of them here.

Make an infinite loop

There are at least three methods to make an infinite loop:

  1. By the count function

2. By the cycle function

3. By the repeat function

Combine multiple iterators into one

The chain() function can help us combine multiple iterators into one.

Pick out adjacent duplicate elements

The groupby function is to pick out adjacent duplicate items in an iterator and put them together.

As shown above, the adjacent same characters were put into together. Furthermore, we can tell the groupby function how to determine two items are the same or not:

Customise a Loop by Ourselves

After enjoying all the above examples, it’s time to think about why the for loops in Python are so flexible and elegant. As far as I am concerned, it’s because we can apply functions into the iterator of a for loop. All the tricks mentioned above are just using some special functions to the iterator. The template of all tricks is as follows:

Under the hood, the built-in itertools module just implements some commonly used functions for us. If we happen to forget a function in it or cannot find the function we need, we can just write one by ourselves. More specifically, these functions are generators. This is why we can generator infinite loops by them.

In a nutshell, we can customise a for loop as we like by writing a customised generator.

Let’s see a simple example:

As the above example shows, we define a generator called even_only. If we use this generator in a for loop, only even numbers will be iterated of the list.

Of course, the above example is just used for explanation purpose. There are other ways to do the same thing, such as using the list comprehension.

Conclusion

Writing loops in Python can be very flexibly and elegantly. We can use some built-in tools properly or even define generators by ourselves to write neat and simple loops.

Thanks for reading. If you like it, don’t forget to follow me to get more great articles about programming and technologies!

Relative articles:

TechToFreedom

Dive into technology, investment and entrepreneurship.

Medium is an open platform where 170 million readers come to find insightful and dynamic thinking. Here, expert and undiscovered voices alike dive into the heart of any topic and bring new ideas to the surface. Learn more

Follow the writers, publications, and topics that matter to you, and you’ll see them on your homepage and in your inbox. Explore

If you have a story to tell, knowledge to share, or a perspective to offer — welcome home. It’s easy and free to post your thinking on any topic. Write on Medium

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store