Math

Number Theory

import math

# Greatest Common Divisor (GCD) - Euclidean algorithm
def gcd(a: int, b: int) -> int:
    while b:
        a, b = b, a % b
    return a

# Using math.gcd
g1 = gcd(48, 18)          # 6
g2 = math.gcd(48, 18)     # 6

# Least Common Multiple (LCM) via GCD
def lcm(a: int, b: int) -> int:
    return a // gcd(a, b) * b

# Simple prime check (good for n up to about 1e7 in worst case)
def is_prime(n: int) -> bool:
    if n < 2:
        return False
    if n % 2 == 0:
        return n == 2
    d = 3
    while d * d <= n:
        if n % d == 0:
            return False
        d += 2
    return True

# Sieve of Eratosthenes - primes up to n
def sieve(n: int) -> list[bool]:
    is_prime = [True] * (n + 1)
    is_prime[0:2] = [False, False]
    p = 2
    while p * p <= n:
        if is_prime[p]:
            for m in range(p * p, n + 1, p):
                is_prime[m] = False
        p += 1
    return is_prime

# Prime factorization (returns list of (prime, exponent))
def factorize(n: int) -> list[tuple[int, int]]:
    factors = []
    d = 2
    while d * d <= n:
        if n % d == 0:
            cnt = 0
            while n % d == 0:
                n //= d
                cnt += 1
            factors.append((d, cnt))
        d += 1
    if n > 1:
        factors.append((n, 1))
    return factors

# Euler's Totient Function φ(n)
def phi(n: int) -> int:
    result = n
    for p, _ in factorize(n):
        result -= result // p
    return result

# Modular arithmetic basics
MOD = 10**9 + 7
a, b = 5, 3
add = (a + b) % MOD
sub = (a - b) % MOD
mul = (a * b) % MOD

# Modular inverse using Fermat's Little Theorem (MOD must be prime)
inv_a = pow(a, MOD - 2, MOD)  # a^{-1} mod MOD

# Number of divisors and sum of divisors from factorization
def divisor_count_and_sum(n: int) -> tuple[int, int]:
    factors = factorize(n)
    cnt = 1
    total_sum = 1
    for p, e in factors:
        cnt *= (e + 1)
        # geometric series: 1 + p + ... + p^e
        total_sum *= (p**(e + 1) - 1) // (p - 1)
    return cnt, total_sum

# Coprime check
def are_coprime(x: int, y: int) -> bool:
    return gcd(x, y) == 1

Combinatorics

import math

# Factorial helper
def fact(n: int) -> int:
    res = 1
    for i in range(2, n + 1):
        res *= i
    return res

# nCr (combinations) using factorial
def nCr(n: int, r: int) -> int:
    if r < 0 or r > n:
        return 0
    return fact(n) // (fact(r) * fact(n - r))

# nPr (permutations)
def nPr(n: int, r: int) -> int:
    if r < 0 or r > n:
        return 0
    res = 1
    for i in range(n, n - r, -1):
        res *= i
    return res

# Python 3.8+ / 3.11+ built-ins (if available)
try:
    c_5_2 = math.comb(5, 2)   # 10
    p_5_2 = math.perm(5, 2)   # 20
except AttributeError:
    c_5_2 = nCr(5, 2)
    p_5_2 = nPr(5, 2)

# Build first k rows of Pascal's Triangle
def pascal(k: int) -> list[list[int]]:
    tri = []
    for i in range(k):
        row = [1] * (i + 1)
        for j in range(1, i):
            row[j] = tri[i - 1][j - 1] + tri[i - 1][j]
        tri.append(row)
    return tri

# Number of subsets of an n-element set
def subset_count(n: int) -> int:
    return 1 << n  # 2**n

# Inclusion–exclusion for |A ∪ B|
def union_two(a: int, b: int, both: int) -> int:
    return a + b - both

# Simple Catalan number (small n)
def catalan(n: int) -> int:
    return nCr(2 * n, n) // (n + 1)

# Ways to choose unordered pairs from n items
def pairs_from_n(n: int) -> int:
    return n * (n - 1) // 2

# Stars and bars: number of non-negative integer solutions to x1+...+xk = n
def stars_and_bars(n: int, k: int) -> int:
    return nCr(n + k - 1, k - 1)

Probability & Expectation

import random

# Basic probability from counts
def probability(successes: int, total: int) -> float:
    return successes / total if total > 0 else 0.0

# Conditional probability P(A|B) = P(A∩B)/P(B)
def conditional(joint: float, p_b: float) -> float:
    return joint / p_b if p_b > 0 else 0.0

# Expected value of discrete variable
def expected_value(values: list[float], probs: list[float]) -> float:
    return sum(v * p for v, p in zip(values, probs))

# Expected number of trials until first success in geometric(p)
def geometric_expectation(p: float) -> float:
    return 1.0 / p

# Linearity of expectation: sum of dice
def expected_sum_of_dice(n: int) -> float:
    # Each die: (1+2+3+4+5+6)/6 = 3.5
    return n * 3.5

# Monte Carlo estimate of probability
def monte_carlo_estimate(trials: int = 10000) -> float:
    hits = 0
    for _ in range(trials):
        x, y = random.random(), random.random()
        if x * x + y * y <= 1:
            hits += 1
    return hits / trials  # ~π/4

# Single step of 1D random walk
def random_walk_step(pos: int) -> int:
    step = random.choice([-1, 1])
    return pos + step

# Normalize weights into probabilities
def normalize(weights: list[float]) -> list[float]:
    s = sum(weights)
    return [w / s for w in weights] if s > 0 else [0.0] * len(weights)

# Build prefix sums for weighted random choice
def build_prefix(probs: list[float]) -> list[float]:
    prefix = []
    s = 0.0
    for p in probs:
        s += p
        prefix.append(s)
    return prefix

def pick_index(prefix: list[float]) -> int:
    r = random.random() * prefix[-1]
    lo, hi = 0, len(prefix) - 1
    while lo < hi:
        mid = (lo + hi) // 2
        if prefix[mid] >= r:
            hi = mid
        else:
            lo = mid + 1
    return lo

Algebra & Series

import math

# Quadratic formula: ax^2 + bx + c = 0
def solve_quadratic(a: float, b: float, c: float) -> tuple[complex, complex]:
    if a == 0:
        # Linear: bx + c = 0
        return (-c / b, )
    disc = b * b - 4 * a * c
    if disc >= 0:
        root_disc = math.sqrt(disc)
    else:
        root_disc = complex(0, math.sqrt(-disc))
    return ((-b - root_disc) / (2 * a), (-b + root_disc) / (2 * a))

# Sum of first n integers: 1 + ... + n
def sum_first_n(n: int) -> int:
    return n * (n + 1) // 2

# Sum of squares 1^2 + ... + n^2
def sum_squares(n: int) -> int:
    return n * (n + 1) * (2 * n + 1) // 6

# Sum of cubes 1^3 + ... + n^3
def sum_cubes(n: int) -> int:
    s = sum_first_n(n)
    return s * s

# Arithmetic progression sum: a, a+d, ..., a+(n-1)d
def ap_sum(a: int, d: int, n: int) -> int:
    return n * (2 * a + (n - 1) * d) // 2

# Geometric progression sum: a + ar + ... + ar^{n-1}, r != 1
def gp_sum(a: float, r: float, n: int) -> float:
    if r == 1:
        return a * n
    return a * (1 - r**n) / (1 - r)

# Clamp a value to [lo, hi]
def clamp(x: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, x))

# Floor and ceil division patterns
def floor_div(a: int, b: int) -> int:
    return a // b

def ceil_div(a: int, b: int) -> int:
    return (a + b - 1) // b

# Safe mid for binary search
def safe_mid(lo: int, hi: int) -> int:
    return lo + (hi - lo) // 2

# Normalize angle to [0, 2π)
def normalize_angle(theta: float) -> float:
    two_pi = 2 * math.pi
    return theta % two_pi

Geometry & Coordinates

import math
from typing import Tuple

Point = Tuple[float, float]

# Euclidean distance between two points
def dist(p1: Point, p2: Point) -> float:
    return math.hypot(p1[0] - p2[0], p1[1] - p2[1])

# Manhattan distance
def manhattan(p1: Point, p2: Point) -> float:
    return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1])

# Cross product of vectors AB x AC
def cross(ax: float, ay: float, bx: float, by: float) -> float:
    return ax * by - ay * bx

def cross_points(a: Point, b: Point, c: Point) -> float:
    return cross(b[0] - a[0], b[1] - a[1], c[0] - a[0], c[1] - a[1])

# Orientation: >0 counterclockwise, <0 clockwise, 0 collinear
def orientation(a: Point, b: Point, c: Point) -> float:
    return cross_points(a, b, c)

# Area of triangle (absolute value)
def triangle_area(a: Point, b: Point, c: Point) -> float:
    return abs(cross_points(a, b, c)) / 2.0

# Polygon area by Shoelace formula
def polygon_area(poly: list[Point]) -> float:
    area = 0.0
    n = len(poly)
    for i in range(n):
        x1, y1 = poly[i]
        x2, y2 = poly[(i + 1) % n]
        area += x1 * y2 - x2 * y1
    return abs(area) / 2.0

# Check if point p lies on segment ab
def on_segment(a: Point, b: Point, p: Point) -> bool:
    if abs(cross_points(a, b, p)) > 1e-9:
        return False
    return (min(a[0], b[0]) <= p[0] <= max(a[0], b[0]) and
            min(a[1], b[1]) <= p[1] <= max(a[1], b[1]))

# Axis-aligned rectangle intersection
def rects_intersect(a: Point, a2: Point, b: Point, b2: Point) -> bool:
    # rectangles: a->a2 and b->b2 as opposite corners
    ax1, ay1 = min(a[0], a2[0]), min(a[1], a2[1])
    ax2, ay2 = max(a[0], a2[0]), max(a[1], a2[1])
    bx1, by1 = min(b[0], b2[0]), min(b[1], b2[1])
    bx2, by2 = max(b[0], b2[0]), max(b[1], b2[1])
    return not (ax2 < bx1 or bx2 < ax1 or ay2 < by1 or by2 < ay1)

# Degrees ↔ radians
def deg_to_rad(deg: float) -> float:
    return math.radians(deg)

def rad_to_deg(rad: float) -> float:
    return math.degrees(rad)

Matrix & Grid

from typing import List, Tuple

Matrix = List[List[int]]

# Matrix multiplication (A: n x m, B: m x p)
def mat_mul(A: Matrix, B: Matrix) -> Matrix:
    n, m, p = len(A), len(A[0]), len(B[0])
    C = [[0] * p for _ in range(n)]
    for i in range(n):
        for k in range(m):
            for j in range(p):
                C[i][j] += A[i][k] * B[k][j]
    return C

# Transpose of matrix
def transpose(A: Matrix) -> Matrix:
    return [list(row) for row in zip(*A)]

# Rotate square matrix 90 degrees clockwise (in-place)
def rotate_90(A: Matrix) -> None:
    n = len(A)
    # transpose
    for i in range(n):
        for j in range(i + 1, n):
            A[i][j], A[j][i] = A[j][i], A[i][j]
    # reverse each row
    for i in range(n):
        A[i].reverse()

# Directions for 4-neighbor grid BFS/DFS
D4: list[Tuple[int, int]] = [(1, 0), (-1, 0), (0, 1), (0, -1)]
D8: list[Tuple[int, int]] = D4 + [(1, 1), (1, -1), (-1, 1), (-1, -1)]

def in_bounds(r: int, c: int, rows: int, cols: int) -> bool:
    return 0 <= r < rows and 0 <= c < cols

# Build 2D prefix sum
def build_prefix_2d(grid: Matrix) -> Matrix:
    rows, cols = len(grid), len(grid[0])
    pref = [[0] * (cols + 1) for _ in range(rows + 1)]
    for r in range(rows):
        for c in range(cols):
            pref[r + 1][c + 1] = (grid[r][c] + pref[r][c + 1] +
                                  pref[r + 1][c] - pref[r][c])
    return pref

# Query sum of submatrix [r1, r2] x [c1, c2] inclusive using prefix
def query_prefix_2d(pref: Matrix, r1: int, c1: int, r2: int, c2: int) -> int:
    return (pref[r2 + 1][c2 + 1] - pref[r1][c2 + 1] -
            pref[r2 + 1][c1] + pref[r1][c1])

# Diagonal traversal (top-left to bottom-right)
def diagonals(grid: Matrix) -> list[list[int]]:
    rows, cols = len(grid), len(grid[0])
    ans: list[list[int]] = []
    for s in range(rows + cols - 1):
        cur = []
        for r in range(rows):
            c = s - r
            if 0 <= c < cols:
                cur.append(grid[r][c])
        ans.append(cur)
    return ans

# Spiral order traversal skeleton
def spiral_order(grid: Matrix) -> list[int]:
    if not grid or not grid[0]:
        return []
    res = []
    top, left = 0, 0
    bottom, right = len(grid) - 1, len(grid[0]) - 1
    while top <= bottom and left <= right:
        for c in range(left, right + 1):
            res.append(grid[top][c])
        top += 1
        for r in range(top, bottom + 1):
            res.append(grid[r][right])
        right -= 1
        if top <= bottom:
            for c in range(right, left - 1, -1):
                res.append(grid[bottom][c])
            bottom -= 1
        if left <= right:
            for r in range(bottom, top - 1, -1):
                res.append(grid[r][left])
            left += 1
    return res

Powers, Logs & Roots

import math

# Binary exponentiation (fast pow)
def fast_pow(base: int, exp: int) -> int:
    res = 1
    b = base
    e = exp
    while e > 0:
        if e & 1:
            res *= b
        b *= b
        e >>= 1
    return res

# Modular exponentiation
def mod_pow(base: int, exp: int, mod: int) -> int:
    return pow(base, exp, mod)

# Check if n is power of two
def is_power_of_two(n: int) -> bool:
    return n > 0 and (n & (n - 1)) == 0

# Integer square root
def int_sqrt(n: int) -> int:
    return math.isqrt(n)

# nth root via binary search (integer root)
def int_nth_root(n: int, k: int) -> int:
    lo, hi = 0, n
    while lo <= hi:
        mid = (lo + hi) // 2
        p = mid**k
        if p == n:
            return mid
        if p < n:
            lo = mid + 1
        else:
            hi = mid - 1
    return hi  # floor root

# Logarithm base change: log_b(x)
def log_base(x: float, b: float) -> float:
    return math.log(x) / math.log(b)

# Highest power of 2 <= n
def highest_power_of_two(n: int) -> int:
    if n <= 0:
        return 0
    return 1 << (n.bit_length() - 1)

# Count bits using bit_length
def count_bits(n: int) -> int:
    return n.bit_length()

# Example: pow with negative exponent (floating result)
negative_exponent = pow(2, -3)   # 0.125

Fractions & Modular Arithmetic

from fractions import Fraction

# Exact rational using Fraction
f1 = Fraction(1, 3)
f2 = Fraction(2, 5)
f_sum = f1 + f2          # 11/15

# Reduce a fraction via GCD
from math import gcd

def reduce_fraction(num: int, den: int) -> tuple[int, int]:
    g = gcd(num, den)
    return num // g, den // g

# Fraction to decimal (no repeat detection)
def fraction_to_decimal(num: int, den: int) -> float:
    return num / den

MOD = 10**9 + 7

# Modular addition, subtraction, multiplication
def mod_add(a: int, b: int, m: int = MOD) -> int:
    return (a + b) % m

def mod_sub(a: int, b: int, m: int = MOD) -> int:
    return (a - b) % m

def mod_mul(a: int, b: int, m: int = MOD) -> int:
    return a * b % m

# Normalize negative modulo result
def normalize_mod(x: int, m: int = MOD) -> int:
    return (x % m + m) % m

# Extended GCD for modular inverse when mod not prime
def extended_gcd(a: int, b: int) -> tuple[int, int, int]:
    if b == 0:
        return a, 1, 0
    g, x1, y1 = extended_gcd(b, a % b)
    return g, y1, x1 - (a // b) * y1

def mod_inv(a: int, m: int) -> int | None:
    g, x, _ = extended_gcd(a, m)
    if g != 1:
        return None
    return x % m

# Modular division a / b mod m
def mod_div(a: int, b: int, m: int = MOD) -> int | None:
    inv_b = mod_inv(b, m)
    if inv_b is None:
        return None
    return a * inv_b % m

Digit & Representation Tricks

# Sum of digits
def digit_sum(n: int) -> int:
    s = 0
    x = abs(n)
    while x:
        s += x % 10
        x //= 10
    return s

# Digital root (repeated digit sum)
def digital_root(n: int) -> int:
    if n == 0:
        return 0
    return 1 + (n - 1) % 9

# Reverse integer
def reverse_int(n: int) -> int:
    sign = -1 if n < 0 else 1
    x = abs(n)
    rev = 0
    while x:
        rev = rev * 10 + x % 10
        x //= 10
    return sign * rev

# Palindrome number check
def is_palindrome_number(n: int) -> bool:
    return n >= 0 and n == reverse_int(n)

# Count digits in base 10
def count_digits(n: int) -> int:
    if n == 0:
        return 1
    c = 0
    x = abs(n)
    while x:
        c += 1
        x //= 10
    return c

# Convert to binary / hex string
def to_binary(n: int) -> str:
    return bin(n)[2:]

def to_hex(n: int) -> str:
    return hex(n)[2:]

# Extract k-th digit from right (0-based)
def kth_digit(n: int, k: int) -> int:
    return abs(n) // (10**k) % 10

# Count set bits using bin
def count_set_bits(n: int) -> int:
    return bin(n).count(\"1\")

# Build number from digits list
def build_number(digits: list[int]) -> int:
    x = 0
    for d in digits:
        x = x * 10 + d
    return x

# Generate all digit masks for n-digit number (0..(1<<n)-1)
def digit_masks(n: int) -> list[int]:
    return list(range(1 << n))