Input

Single integer input

x = int(input())

Single string input

x = input()

Line with multiple integers

# one
x, y, z = map(int, input().split())
# two
xyz = [int(t) for t in input().split()]

Handling large inputs

If the input data is large, input() might be slow. You can use stdin.readline() from the sys module to speed up input processing. stdin.readline() reads the input as a string, and you need to cast it to int or float as needed.

⚠️Note:

stdin.readline() also reads the trailing \n character, which can be removed using .strip() or slicing.

from sys import stdin
n = int(input())
for _ in range(n):
  x, y, z = map(int, stdin.readline().split())
# or 
	xyz = [list(map(int, stdin.readline().split())) for _ in range(n)]

Multiple lines of input, each line containing a string

from sys import stdin
for _ in range(n):
  s = stdin.readline()[:-1]
# or back list [[], [], [], ...[]]
s = [stdin.readline()[:-1] for _ in range(n)]

Output

Outputting YES or NO

When condition is True, the first character is N, and taking every second character gives No. Conversely, if condition is False, the first character is Y, resulting in Yes.

print("YNeos"[condition::2])

You can simplify the code using slicing.

n = input()
if (n == '3') or (n == '5') or (n == '7'):
    print('YES')
else:
    print('NO')

print('NYOE S'[input()in'357'::2])

Lists (Arrays)

One-dimensional list

x = [0 for _ in range(n)]
# or
x = [0] * n

Two-dimensional list

x = [[0] * m for _ in range(n)]

Flattening a nested list

To flatten a nested list, use an outer loop to iterate over each sublist and an inner loop to iterate over each element within those sublists.

data = [[1,2,3],[4,5,6],[7,8,9]]
L = [x for y in data for x in y]
print(b)

[1,2,3,4,5,6,7,8,9]

Sorting a two-dimensional list

Sort by the second element of each sublist.

L = [[1, 4, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6], [2, 3, 4], [1, 5, 3], [2, 3, 4], [5, 6, 7]]
L = sorted(L, key=lambda x: x[1])
print(L)

[[2, 3, 4], [2, 3, 4], [2, 3, 4], [1, 4, 3], [3, 4, 5], [4, 5, 6], [1, 5, 3], [5, 6, 7]]

String

String concatenation

Directly appending characters to a string is inefficient with a time complexity of $\mathcal{O}(N^2)$. Instead, append to a list and use join to concatenate the list elements, which has a time complexity of $\mathcal{O}(N \log N)$.

last = []
for x in s:
  if condition:
    last.append(x)
res = ''.join(last)

Loop Structures

Accelerating loops with lists

When the index value is not needed, avoid using range and use the following approach instead:

N = 10
for _ in [0] * N:
  ....

You can also optimize nested loops by pre-computing values in a list to reduce redundant calculations and make better use of caching:

for i in range(X):
  tmp_Y = list(range(Y))
  for j in tmp_Y:
    .....

Improving Speed

Global variables vs. local functions

Global variables are slightly slower. Consider placing code inside a main function and calling it:

from sys import stdin
def main():
  from builtins import input, in...
  ....
main()

Faster stdin.readline()

Assigning stdin.readline() to a variable can improve speed:

from sys import stdin
def main():
  readline = stdin.readline
  a, b = map(int, readline().split())

Infinity

Assign float('inf') to a variable for representing infinity:

INF = float('inf')

Miscellaneous

collections Library

Counter is a dictionary-like class that provides methods for counting hashable objects and supports addition and subtraction of counts.

from collections import Counter
L1 = [1,2,3,4,5]
L2 = [1,2,3,4,5,6]

a = Counter(L1)
b = Counter(L2)

print(a+b)

print(b-a)

a = a + b
print(a.items())

print(a.keys())

print(a.values())

for k, v in a.items():
  print(k, v)

Counter({1: 2, 2: 2, 3: 2, 4: 2, 5: 2, 6: 1})
Counter({6: 1})
dict_items([(1, 2), (2, 2), (3, 2), (4, 2), (5, 2), (6, 1)])
dict_keys([1, 2, 3, 4, 5, 6])
dict_values([2, 2, 2, 2, 2, 1])
1 2
2 2
3 2
4 2
5 2
6 1

Math Library

Integer square root

from math import isqrt
isqrt(5)

2

To round up, use: 1 + isqrt(x-1)

from math import isqrt
1 + isqrt(5-1)

Least common multiple (LCM) and greatest common divisor (GCD)

Python 3.9+ supports an arbitrary number of arguments

from math import lcm, gcd
gcd(10, 5)
lcm(10, 5)

5
10

itertools Library

accumulate

Computes cumulative sums and returns them as a sequence:

from itertools import*
L = [1,2,3]
print(list(accumulate(L)))

[1, 3, 6]

groupby

Groups consecutive identical elements, separating non-consecutive elements into distinct groups. It can also be used with lambda expressions for custom grouping.

from itertools import groupby

L = [1, 1, 2, 3, 3, 3, 1, 2, 2]

for key, value in groupby(L):
    print(key, list(value))

1 [1, 1]
2 [2]
3 [3, 3, 3]
1 [1]
2 [2, 2]

from itertools import groupby

a = [1, 3, 2, 4, 3, 1, 1, 2, 4]

for key, value in groupby(a, key=lambda x: x % 2):
    print(key, list(value))

1 [1, 3]
0 [2, 4]
1 [3, 1, 1]
0 [2, 4]

Base Conversion

Use the int() function with a base parameter to convert to decimal:

Strings with prefixes starting with 0b denote binary, while those starting with 0x denote hexadecimal

print(int("1010", 2))
print(int("1A", 16))
print(int('0b1010'), 2)
print(int("0xAC"), 16)

10
26
10
172

Use bin() for binary and hex() for hexadecimal conversions:

print(bin(255)) 
print(hex(255))

0b11111111
0xff