Printing 1,000,000 numbers in 1 sec. in Python

Question

Here's a fairly simple task from CSES Problem Set - Permutations 1070 that reads:

A permutation of integers 1,2, …, n is called beautiful if there are no adjacent elements whose difference is 1.

Given n, construct a beautiful permutation if such a permutation exist

The constraints are pretty tight:

• Time limit: 1.00 s
• Memory limit: 512 MB
• 1 ≤ n ≤ 10^6

Here's the code:

n = int(input())
if n == 2 or n == 3:
    print("NO SOLUTION")
elif n == 1:
    print(1)
elif n == 4:
    print("3 1 4 2")
else:
    for i in range(1, n + 1, 2):
        print(str(i) + " ", end=" ")
    for i in range(2, n + 1, 2):
        print(str(i) + " ", end=" ")

It passes all tests except for n = 1000000. For that test it takes 15 sec. on my machine. The code is in Python 3.8

The question is, what can can be improved in terms of printing the numbers?

Answer 1

Nice solution, few suggestions:

• Printing the numbers one by one might be the issue. Generate the list of numbers first and then call print only once.
• The case n==1 is already handled, so the first elif can be removed.

Applying the suggestions:

n = int(input())
if n == 2 or n == 3:
    print("NO SOLUTION")
elif n == 4:
    print("3 1 4 2")
else:
    beautiful_perm = [*range(1, n + 1, 2), *range(2, n + 1, 2)]
    print(' '.join(map(str, beautiful_perm)))

By inverting the ranges we don't need to check for n==4:

n = int(input())
if n == 2 or n == 3:
    print("NO SOLUTION")
else:
    beautiful_perm = [*range(2, n + 1, 2), *range(1, n + 1, 2)]
    print(' '.join(map(str, beautiful_perm)))

Runtime on CSES:

n = 906819 (CPython3)
Original: 0.92 s
Improved: 0.26 s

n = 1000000 (CPython3)
Original: timeout
Improved: 0.28 s

n = 1000000 (PyPy3)
Original: 0.61 s
Improved: 0.15 s

Answer 2

Your code gets accepted as-is. Just need to choose PyPy3 instead of CPython3.

Another version that also gets accepted with CPython3 (using Marc's logic but a simpler way to print):

n = int(input())
if 2 <= n <= 3:
    print("NO SOLUTION")
else:
    print(*range(2, n + 1, 2), *range(1, n + 1, 2))

Printing like that moves the loop from your own Python loop with lots of print calls to a single call and looping in C. Even faster is Marc's way with ' '.join, though. Times for test case #20, where n=906819, the largest where yours is fast enough to not get killed:

        CPython3  PyPy3
Yours   0.93 s    0.56 s
Mine    0.37 s    0.32 s
Marc's  0.25 s    0.14 s

Why is Marc's way even faster? The print documentation says "The file argument must be an object with a write(string) method". And if we use our own such object we can see a lot fewer write calls in Marc's than in mine:

class Write:
    def write(self, string):
        print(f'write({string!r})')

nums = range(5)

print(' '.join(map(str, nums)), file=Write())
print()
print(*nums, file=Write())

Output:

write('0 1 2 3 4')
write('\n')

write('0')
write(' ')
write('1')
write(' ')
write('2')
write(' ')
write('3')
write(' ')
write('4')
write('\n')

The overhead of all those function calls is probably what makes mine slower. The overhead of ' '.join(... just seems to be smaller.

That said, I usually prefer my way if it's fast enough and thus try that first. Except maybe in a competition where I have reason to believe it might not be fast enough and there's a penalty for unsuccessful submissions.