Python Standard Library для спортивного программирования / Хабр

sys

Ввод через input() относительно медленный. Причина - операции обработки, преобразования и проверка окончания строки. В задачах с большими входными данными это может привести к Time Limit Exceeded. Поэтому используется быстрый ввод черезsys.stdin.readline.

Быстрое чтение n чисел:

import sys

data = sys.stdin.read().split()
n = int(data[0])

arr = list(map(int, data[1:]))

collections

Документация по модулю collections

Содержит структуры данных, удобные для алгоритмических задач.

deque

Двусторонняя очередь с операциями O(1).

from collections import deque

q = deque()

q.append(1)
q.append(2)

print(q.popleft())

Использование: BFS, очереди, симуляции, sliding window.

from collections import deque

def bfs(start, graph):
    q = deque([start])
    visited = {start}

    while q:
        v = q.popleft()

        for to in graph[v]:
            if to not in visited:
                visited.add(to)
                q.append(to)

Counter

Подсчёт количества элементов.

from collections import Counter

arr = [1, 1, 2, 3, 3, 3]

cnt = Counter(arr)

print(cnt[3])

Использование: частоты, анаграммы, multiset.

defaultdict

Словарь с автоматически создаваемым значением.

from collections import defaultdict

g = defaultdict(list)

g[1].append(2)
g[1].append(3)

from collections import defaultdict

g = defaultdict(list)

edges = [(1, 2),(1, 3),(2, 4)]

for u,v in edges:
    g[u].append(v)
    g[v].append(u)

heapq

Документация по модулю heapq

Реализация min-heap (куча).

Использование: алгоритм Дейкстры, поиск k минимальных элементов, задачи с приоритетами.

import heapq

h = []
heapq.heappush(h, 5)
heapq.heappush(h, 2)
heapq.heappush(h, 8)

print(heapq.heappop(h))  # 2

import heapq

arr = [7, 2, 5, 1, 9]

heapq.heapify(arr)

for _ in range(3):
    print(heapq.heappop(arr))

import heapq

def dijkstra(start, graph):

    dist = {v: float("inf") for v in graph}
    dist[start] = 0

    pq = [(0, start)]

    while pq:

        d,v = heapq.heappop(pq)

        if d > dist[v]:
            continue

        for to,w in graph[v]:

            nd = d + w

            if nd < dist[to]:
                dist[to] = nd
                heapq.heappush(pq, (nd, to))

    return dist

bisect

Документация по модулю bisect

Модуль для бинарного поиска в отсортирова��ном массиве.

Основные функции: bisect_left, bisect_right, insort.

import bisect

arr = [1, 3, 5, 7]
pos = bisect.bisect_left(arr, 4)
print(pos)

Вставка с сохранением сортировки

import bisect

arr = [1, 3, 5, 7]
bisect.insort(arr, 4)
print(arr)

itertools

Документация по модулю itertools

Генераторы комбинаторных объектов.

permutations

from itertools import permutations

for p in permutations([1, 2, 3]):
    print(p)

Использование: bruteforce, перебор порядков.

combinations

from itertools import combinations

for c in combinations([1, 2, 3], 2):
    print(c)

Использование: выбор k из n, перебор подмножеств.

product

Декартово произведение.

from itertools import product

for p in product([1, 2], [3, 4]):
    print(p)

accumulate

Префиксные суммы.

from itertools import accumulate

arr = [1, 2, 3, 4]
pref = list(accumulate(arr))

functools

Документация по модулю functools

Полезные инструменты функционального программирования.

lru_cache

Мемоизация результатов функции.

from functools import lru_cache

@lru_cache(None)
def fib(n):

    if n < 2:
        return n

    return fib(n - 1) + fib(n - 2)

Использование: DP, рекурсия, overlapping subproblems.

math

Документация по модулю math

Наиболее полезные математические функции: gcd, lcm, sqrt, factorial, comb, log, ceil, floor.

gcd

import math

print(math.gcd(12, 18))

factorial

import math

print(math.factorial(5))

Комбинации

import math

print(math.comb(5, 2))

operator

Документация по модулю operator

Функциональные версии операторов. Часто используются в сортировке.

itemgetter

from operator import itemgetter

arr = [(1, 5), (2, 3), (3, 1)]
arr.sort(key=itemgetter(1))
print(arr)

array

Документация по модулю array

Компактный массив чисел. Используется когда важна память.

from array import array

a = array('i', [1, 2, 3, 4])

statistics

Документация по модулю statistics

Статистические функции.

import statistics

arr = [1, 2, 3, 4]

print(statistics.mean(arr))

random

Документация по модулю random

Используется для генерации случайных данных.

import random

print(random.randint(1, 10))

Полезно знать - PyPy

Если решение на Python укладывается в алгоритмические ограничения, но не проходит по времени, иногда достаточно запустить тот же код под PyPy.
JIT-компилятор (Just-In-Time) анализирует код во время выполнения, оптимизирует часто выполняемые участки и компилирует их в машинный код. За счёт этого PyPy нередко работает быстрее в задачах, где значительная часть времени уходит на повторяющиеся вычисления.
Для коротких программ выигрыш может отсутствовать из-за времени «разогрева» JIT-компилятора, кроме того, PyPy обычно потребляет больше памяти.

Если заметите опечатки, неточности или у вас есть идеи, как можно дополнить материал - напишите об этом в комментариях.