Some easy but useful optimization in Python.

Revision en4, by Dart-Xeyter, 2020-10-31 22:25:09

I have a few friends who write rounds in Python, and I noticed that they don't use some very simple optimizations, and the program ends up getting TL. Meanwhile, if you use these constructions, in half or more cases TL is removed.

I will show everything using the example code of my friend I_am_Drew from a recent round. This code received TL (worked longer than $$$1$$$ second).

for __ in range(int(input())):
    n, m = list(map(int, input().split()))
    kek = []
    for i in range(n):
        el = list(map(int, input().split()))
        el.append(0)
        kek.append(el)
    stolb = list(map(int, input().split()))
    ind = 0
    for i in range(m):
        if kek[0][i] in stolb:
            ind = i
            break
 
    for i in range(n):
        kek[i][m] = stolb.index(kek[i][ind])
 
    for j in range(m-1):
        stolb = list(map(int, input().split()))
 
    kek.sort(key=lambda x: x[m])
    for elem in kek:
        elem.pop()
        print(*elem)

First, as you know, data input and output takes quite a long time. Fortunately, this can be fixed using the sys module. I usually write this way because it's the quickest fix, but of course it's not exactly code-style :)

from sys import stdin, stdout
input, print = stdin.readline, stdout.write

The stdin.readline function reads a string like input, but faster. Also, if necessary, there is, for example, the stdin.read function, which reads all input as a string (then you need to put ^D after it is completed), and others, but I usually do not use them. The output is more complicated, the stdout.write function accepts only strings, and does not output a line feed or other separator after it. Therefore, you have to write as in the example below, it is also not very long to fix it, the main thing is not to forget :) After the conversions, you get this code. (Note that the input code has not changed at all, but the output at the end is quite strong).

from sys import stdin, stdout
input, print = stdin.readline, stdout.write
for __ in range(int(input())):
    n, m = list(map(int, input().split()))
    kek = []
    for i in range(n):
        el = list(map(int, input().split()))
        el.append(0)
        kek.append(el)
    stolb = list(map(int, input().split()))
    ind = 0
    for i in range(m):
        if kek[0][i] in stolb:
            ind = i
            break

    for i in range(n):
        kek[i][m] = stolb.index(kek[i][ind])

    for j in range(m - 1):
        stolb = list(map(int, input().split()))

    kek.sort(key=lambda x: x[m])
    for elem in kek:
        elem.pop()
        # print(' '.join(map(str, elem)))
        for q in elem:
            print(str(q)+' ')
        print('\n')

It is also known that global variables work longer than local ones, so if you put all the code (of course, without other functions) in main, it will also work faster. The final version of the code looks like this:

from sys import stdin, stdout
input, print = stdin.readline, stdout.write


def main():
    for __ in range(int(input())):
        n, m = list(map(int, input().split()))
        kek = [list(map(int, input().split()))+[0] for _ in range(n)]
        stolb = list(map(int, input().split()))
        ind = 0
        for i in range(m):
            if kek[0][i] in stolb:
                ind = i
                break

        for i in range(n):
            kek[i][m] = stolb.index(kek[i][ind])

        for j in range(m - 1):
            stolb = list(map(int, input().split()))

        kek.sort(key=lambda x: x[m])
        for elem in kek:
            print(' '.join(map(str, elem[:-1])))
            print('\n')


main()

Note that there were very few changes, but the program accelerated at least 2 times and now gets OK, working in 545 milliseconds.

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
en8 English Dart-Xeyter 2020-10-31 22:48:44 9 Tiny change: ' is quite strong).\n\n~~~~' -> ' is quite a lot).\n\n~~~~' (published)
en7 English Dart-Xeyter 2020-10-31 22:40:40 23 Tiny change: '1] from a recent round. This cod' -> '1] from a [problem:1413B]. This cod'
en6 English Dart-Xeyter 2020-10-31 22:36:43 33
en5 English Dart-Xeyter 2020-10-31 22:34:39 907
en4 English Dart-Xeyter 2020-10-31 22:25:09 1204
en3 English Dart-Xeyter 2020-10-31 22:07:55 1158
en2 English Dart-Xeyter 2020-10-31 21:53:34 1349
en1 English Dart-Xeyter 2020-10-31 21:24:14 285 Initial revision (saved to drafts)