Introduction
Hello everyone, I have recieved some messages asking how to improve python code to avoid TLE. There have been blogs before but they are rather old, so I will share some tips, hope it helps! These are the tips that I have learnt over the years and some of them may not be well known to non-python mains.
Do let me know in the comments if you have more tips as well.
Essentials (Most questions can be solved using these tips)
Basically, try to make the code as "Pythonic" as possible, the more python gets to call its underlying C functions the faster and better it is.
- Use PyPy 3 64 bit, it is faster than python 3.
- For fast I/O, I use
input = sys.stdin.readline
, then for output, if the output has multiple lines, I just create a list to store the outputs, thenprint("\n".join(lst))
on the list. Because multiple print statements are slow in python. - Put all your code into functions like a main function, as local variables lookup in python is faster than global variables, I think your code will run around ~30% faster just by doing this.
- List creation:
[0]*n
is fastest, list comprehension is also fast. For 2d lists: I do:[[0]*m for _ in range(n)]
, as[[0]*m]*n
will not work due to reference issues. Also,.extend()
is faster than.append()
if you are adding multiple elements to the list at a time. - Use
for i in arr
instead offor i in range(n): arr[i]
as it is faster. If you need the index as well as the data, useenumerate()
as it is faster. Alsofor x, y in arr:
to iterate over multiple elements. - Use built-in methods like
map()
andsum()
as they are faster than implementing the stuff yourself. The more Pythonic the better. - For deep recursion, use an explicit stack instead of a recursive function. Recursion in python is slow and memory intensive.
- Try to keep your code as simple as possible especially in the for loops, to reduce the amount of computation needed and also so that PyPy can optimise your code easily.
- String concatenation in python is slow in a for loop as new intermediate strings are always being created. To solve this, use a list and
.append()
(orextend()
), at the end call"".join(lst)
- Try to not change the data type of variables and collections. Collections should store a single data type. For example PyPy has special dict strategies to speed up dicts with keys of all same types. I believe for lists and other collections there will be some optimisations as well, as python can predict the type better and reduce overhead from mutating types.
🦀-Python
As mentioned in this comment: https://codeforces.net/blog/entry/106541?#comment-949024
Non-Pythonic ways to improve python performance. You won't need to use some of these unless the time limits are tight.
- Flattening 2d lists into 1d: This TLE: https://cses.fi/paste/f6db6ab935e23d2072b133/ This passes: https://cses.fi/paste/c7fa7fcb039aa28f72b17f/
- Packing tuple into ints to sort, then unpacking as sorting tuples are slow: This TLE: https://cses.fi/paste/972de727f25b8f1171102a/ This passes: https://cses.fi/paste/25ce406e0f776551711046/
- Iterating over multiple elements using zip: This TLE: https://cses.fi/paste/1c4afcf48c0400b071e0f8/ This passes: https://cses.fi/paste/85c4f40b683a6eec71e10d/
- In some cases deque() can be faster than list, at least for some questions I've done in CSES graphs.
Other useful tips:
- Learn your standard library properly (take a look at the documentations), Python has some very useful functions and classes, for example I sometimes use Decimal class to deal with floats of higher precision (but take note the precision cannot be too high else TLE). Counter and bisect are nice to use as well.
- Pyrival is good, has lots of algorithms and data structures, lots of useful stuff there like SortedList() and prime factors.
- To avoid integer hash hacking TLE for
set()
anddict()
, you canrandom.shuffle()
the array if order is not important, otherwise you can hash the str(int) or use random int xor to wrap the ints.