Python for Competitive Programming¶

In one sentence

Every solution has the same shape: read the input line by line, compute, and print each answer. This page collects the Python you need for that, and the traps that cost points.

1. What problem does it solve?¶

Example — 2191A · Array Coloring (Raymond's first Accepted, 22 Jan)

\(n\) cards lie in a row, card \(i\) shows \(a_i\), and all \(a_i\) are different (each between \(1\) and \(n\)). Colour every card red or blue so that

any two neighbouring cards in the row have different colours, and
after sorting the cards by their numbers, any two neighbours in sorted order also have different colours.

Is that possible? Limits: \(t \le 200\), \(n \le 100\).

Input          Output
4
4
2 3 4 1        YES
3
2 3 1          NO
5
3 4 1 2 5      YES
5
3 1 4 2 5      NO

Thinking. Rule 1 forces the colours to alternate along the row, so a card's colour is decided by whether its position \(i\) is odd or even. Rule 2 forces them to alternate in sorted order. Since the numbers are exactly \(1..n\), the card showing \(a_i\) is at sorted position \(a_i\), so its colour is decided by whether \(a_i\) is odd or even. Both can hold only if \(i\) and \(a_i\) have the same parity relationship for every card: \((i - a_i) \bmod 2\) is the same for all \(i\).

First test, 2 3 4 1: \(i - a_i = -1, -1, -1, 3\), all odd. YES. Second test, 2 3 1: \(-1, -1, 2\), mixed. NO.

The code needs only the basics: read \(t\) test cases, read a list, loop, test parity with %, and print. All of that is below.

2. The skeleton¶

Every lesson file starts from this template. Here it is from July 19:

import sys

input = sys.stdin.readline


# sys.setrecursionlimit(300000)


def solve():
    n = int(input())

    return


t = 1
# t = int(input())
for _ in range(t):
    solve()

line	why
`input = sys.stdin.readline`	the built-in `input()` is slow for \(10^5\) lines; `readline` is much faster
`# sys.setrecursionlimit(300000)`	uncomment only for deep recursion (see DFS)
`def solve():`	one test case; local variables inside a function are also faster
`t = 1` / `# t = int(input())`	switch to the second line when the input starts with the number of test cases

3. Reading input¶

input line	code
`5`	`n = int(input())`
`3 7`	`n, k = map(int, input().split())`
`3 2 5`	`a = list(map(int, input().split()))`
`#..#` (a string / grid row)	`s = input().strip()`
\(n\) grid rows	`grid = [input().strip() for _ in range(n)]`
`1.5 2.5`	`x, y = map(float, input().split())`

readline keeps the newline

sys.stdin.readline() returns "abc\n", not "abc". For numbers it does not matter, because int() and split() ignore it. For strings, always .strip(). Our August 2 lesson showed this:

def solve():
    # input string
    s = input()
    t = input().strip()

    print(len(s), len(t))

    if s == t:
        print("YES")
    else:
        print("NO")

With input abc and abc, it prints 4 3 and then NO: s still has the "\n".

4. Printing output¶

want	code	prints
one value per line	`print(x)`	`5`
several values, one line	`print(a, b)`	`5 7`
values from a loop, one line	`print(x, end=" ")` inside the loop, then `print()`	`3 10 5 …`
no separator	`print(ch, end="")` inside the loop, then `print()`	`LDDR`
a formatted float	`print(f"{x:.6f}")`	`3.141593`

5. Numbers¶

operation	example	result	note
integer division	`7 // 2`	`3`	rounds down: `-7 // 2` is `-4`
true division	`7 / 2`	`3.5`	always a `float`; avoid for large integers
remainder	`7 % 2`	`1`	the result has the sign of the divisor: `-7 % 2` is `1`
power	`2 ** 10`	`1024`
big integers	`10 ** 30`	exact	Python `int` never overflows
modulo arithmetic	`(a * b) % MOD`		apply `%` after every step in counting problems

n / 2 where you meant n // 2

/ gives a float. For \(n\) around \(10^{17}\) a float cannot represent every integer, so answers silently go wrong. Use // for integers.

6. Lists and 2D lists¶

Our June 28 lesson checked how to build a 2D list safely:

def solve():
    a = [[0] * 3 for _ in range(3)]

    print(a)

    a[0][0] = 1

    print(a)

It prints [[0, 0, 0], [0, 0, 0], [0, 0, 0]] and then [[1, 0, 0], [0, 0, 0], [0, 0, 0]]: only row \(0\) changed.

The tempting shortcut a = [[0] * 3] * 3 is wrong. It makes one row and three references to it, so a[0][0] = 1 changes every row to [1, 0, 0]. Always use the list comprehension.

7. The containers you will use¶

container	create	main operations	cost	used for
`list`	`[0] * n`	`a[i]`, `append`, `pop()`	\(O(1)\)	arrays, stacks, adjacency lists
`set`	`set()`	`add`, `remove`, `x in s`	\(O(1)\) avg	"have I seen this?", deleting edges (lesson notes)
`dict`	`{}`	`d[k]`, `d.get(k, 0)`, `k in d`	\(O(1)\) avg	counting, mapping values to indices
`collections.Counter`	`Counter(a)`	`c[x]` (missing → `0`)	\(O(1)\) avg	frequency counts
`collections.defaultdict`	`defaultdict(list)`	`d[k].append(v)`	\(O(1)\) avg	groups, graphs with unusual labels
`collections.deque`	`deque()`	`append`, `popleft`, `appendleft`	\(O(1)\)	the BFS queue
`heapq` on a list	`q = []`	`heappush(q, x)`, `heappop(q)`	\(O(\log n)\)	smallest-first: Dijkstra
`bisect` on a sorted list		`bisect_left(a, x)`	\(O(\log n)\)	counting ≤ / < x

heapq is a min-heap. For a max-heap, push -x and negate when popping.

Sorting, including key=lambda, is on Sorting & Greedy.

8. Worked solution — Array Coloring¶

The skeleton from §2 with t = int(input()) switched on. Positions are \(1..n\) in the statement; in Python a[i] has index \(i = 0..n-1\), so the position is i + 1.

import sys

input = sys.stdin.readline


def solve():
    n = int(input())
    a = list(map(int, input().split()))

    first = (1 - a[0]) % 2
    ok = True
    for i in range(n):
        if (i + 1 - a[i]) % 2 != first:
            ok = False

    if ok:
        print("YES")
    else:
        print("NO")


t = int(input())
for _ in range(t):
    solve()

% in Python never returns a negative number (-1 % 2 is 1), so (i + 1 - a[i]) % 2 is always \(0\) or \(1\) and the comparison is safe.

9. Common mistakes¶

Using the built-in input() on big inputs

Reading \(2 \cdot 10^5\) lines with the built-in input() can take longer than the whole algorithm. Put input = sys.stdin.readline at the top.

Forgetting t = int(input())

Most Codeforces problems start with the number of test cases. The skeleton has t = 1 by default, so switch lines.

Comparing strings that still have \n

See §3. .strip() every string you read.

Shadowing built-ins

sum = 0, max = a[0], list = [...] overwrite Python's functions, and a later sum(a) crashes. Pick other names (total, best, arr).

[[0] * m] * n

Every row is the same list (§6).

State left over between test cases

If solve() runs \(t\) times, create lists inside solve(), sized by this test's \(n\).

10. Practice¶

Problem	Where	Point
2191A · Array Coloring	Codeforces, Jan 22	this page
2191B · MEX Reordering	Codeforces, Jan 23 (WA, then AC)	counting, then careful output
2227A · Koshary	Codeforces, May 22	basic I/O
2227B · Party Monster	Codeforces, May 22	reading strings
1742C · Stripes	4-Misc · I	read a grid with `strip()`
1791C · Prepend and Append	3-TwoPointers · A	string indexing

Credits & licenses

§2, §3 and §6 quote our lesson code. Everything else (the Array Coloring analysis, tables, mistakes) is ours.