4 performance optimization tips for faster Python code

Python is one of the most popular programming environments and many developers look to HackerEarth as a way to learn more about the language. Consequently, for the challenges that are hosted on our platform, a huge number of submissions are in Python.

We all know how we are able to pass the initial test case, but the other test cases fail due to memory or time limits. So, let us look at some of the tips that one should keep in mind so that a correct python program remains within the platform’s constraints for the challenges.

1. Get the whole setup ready before-hand

This is common sense. Get the whole setup ready. Fire-up your python editor before-hand. If you are writing your files in your local, create a virtual environment and activate it. Along with this, I would advise one other thing which might seem a bit controversial and counter-intuitive, and that is to use TDD. Use your favourite testing tool. I generally use pytest and have that “pip”-d in my virtual environment and start writing small test scripts. I have found that testing helps in clarity of thought, which helps in writing faster programs. Also, this helps in refactoring the code to make it faster. We will get to it later.

2. Get the code working first

People have their own coding styles. Use the coding style that you are most comfortable with. For the first iteration, make the code work, at least and make the submission. See if it passes for all the test cases. If it’s passing then, c’est fait. It’s done. And you can move on to the next question.

In case its passing for some of the test cases, while failing for others, citing memory issues, then you know that there is still some work left.

3. Python Coding Tips

We can use some pythonic code constructs that give us better performance. Let’s look at them below.

Strings:

Do not use the below construct.

s = ""
for x in somelist:
   s += some_function(x)

Instead use this

slist = [some_function(el) for el in somelist]
s = "".join(slist)

This is because in Python, str is immutable, so the left and right strings have to be copied into the new string for every pair of concatenation.

Language Constructs:

Functions: Make functions of your code and although procedural code is supported in Python, it’s better to write them in functions. You might have heard this advice from proponents of code reusability and champions of TDD (Note: I am one of them), but here I would say that functions are better because they are faster. So

def main():
    for i in xrange(10**8):
        pass

main()

is better than

for i in xrange(10**8):
    pass

This is because of the underlying CPython implementation. In short, it is faster to store local variables than globals. Read more in this SO post.

I would recommend you keep your procedural code as little as possible. You can use the following standard template.

def solution(args):
    # write the code
    pass


def main():
    # write the input logic to take the input from STDIN
    Input_args = ""
    solution(input_args)


If __name__ == "__main__":
    main()

Use the standard library:

Use built-in functions and the standard library as much as possible. Therefore, instead of this:

newlist = []
for item in oldlist:
    newlist.append(myfunc(item))

Use this:

newlist = map(myfunc, oldlist)

There is also the list expressions or the generator expressions.

newlist = [myfunc(item) for item in oldlist] # list expression
newlist = (myfunc(item) for item in oldlist) # generator expression

Similarly, use the standard library, like itertools, as they are generally faster and optimized for common operations. So you can have something like permutation for a loop in just three lines of code.

>>> import itertools
>>> iter = itertools.permutations([1,2,3])
>>> list(iter)
[(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]

But if the containers are tiny, then any difference between code using libraries is likely to be minimal, and the cost of creating the containers will outweigh the gains they give.

Generators:

A Python generator is a function which returns a generator iterator (just an object we can iterate over) by calling yield. When a generator function calls yield, the “state” of the generator function is frozen; the values of all variables are saved and the next line of code to be executed is recorded until next() is called again. Generators are excellent constructs to reduce both the average time complexity as well as the memory footprint of the code that you have written. Just look at the following code for prime numbers.

def fib():
    a, b = 0, 1
    while 1:
        yield a
        a, b = b, a + b

So, this will keep on generating Fibonacci numbers infinitely without the need to keep all the numbers in a list or any other construct. Please keep in mind that you should use this construct only when you don’t have any absolute need to keep all the generated values because then you will lose the advantage of having a generator construct.

4. Algorithms and Data structures

To make your code run faster, the most important thing that you can do is to take two minutes before writing any code and think about the data-structure that you are going to use. Look at the time complexity for the basic python data-structures and use them based on the operation that is most used in your code. The time complexity of the list taken from python wiki is shown below.

Similarly, keep on reading from all sources about the most efficient data structures and algorithms that you can use. Keep an inventory of the common data structures such as nodes and graphs and remember or keep a handy journal on the situations where they are most appropriate.

Writing fast code is a habit and a skill, which needs to be honed over the years. There are no shortcuts. So do your best and best of luck.