Tan Kim

수학 알고리즘

소수 (Prime)

# 에라토스테네스의 체 — n 이하 소수 O(n log log n)
def sieve(n):
    is_prime = [True] * (n+1)
    is_prime[0] = is_prime[1] = False
    for i in range(2, int(n**0.5) + 1):
        if is_prime[i]:
            for j in range(i*i, n+1, i):
                is_prime[j] = False
    return [i for i in range(n+1) if is_prime[i]]
 
# 소수 판별 O(√n)
def is_prime(n):
    if n < 2: return False
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0: return False
    return True
 
# 소인수 분해
def factorize(n):
    factors = {}
    d = 2
    while d * d <= n:
        while n % d == 0:
            factors[d] = factors.get(d, 0) + 1
            n //= d
        d += 1
    if n > 1: factors[n] = 1
    return factors

GCD / LCM

import math
 
math.gcd(12, 8)   # 4
math.lcm(4, 6)    # 12  (Python 3.9+)
 
# 수동 구현
def gcd(a, b):
    while b:
        a, b = b, a % b
    return a
 
def lcm(a, b):
    return a * b // gcd(a, b)

모듈러 산술

MOD = 10**9 + 7
 
# 모듈러 거듭제곱 O(log exp)
pow(base, exp, MOD)  # 내장 함수 사용 권장
 
# 모듈러 역원 (MOD가 소수일 때)
def mod_inv(a, mod):
    return pow(a, mod - 2, mod)  # 페르마 소정리
 
# 조합 C(n, r) mod p
def comb_mod(n, r, mod):
    if r > n: return 0
    num = den = 1
    for i in range(r):
        num = num * (n - i) % mod
        den = den * (i + 1) % mod
    return num * pow(den, mod - 2, mod) % mod

비트 연산

n & (n-1)       # 가장 낮은 비트 제거 (2의 거듭제곱 판별: n & (n-1) == 0)
n & (-n)        # 가장 낮은 비트만 추출
bin(n).count('1')  # 1의 개수 (popcount)
n >> k          # n // 2^k
n << k          # n * 2^k
n ^ n           # 0 (자기 자신과 XOR)
a ^ b ^ b       # a (XOR로 암호화/복호화)
 
# 부분집합 열거
def subsets(n):
    for mask in range(1 << n):
        subset = [i for i in range(n) if mask & (1 << i)]
        print(subset)

투 섬 / 수학 응용

# 빠른 거듭제곱
def fast_pow(base, exp):
    result = 1
    while exp:
        if exp & 1:
            result *= base
        base *= base
        exp >>= 1
    return result
 
# 행렬 거듭제곱 (피보나치 O(log n))
def mat_mul(A, B, mod):
    n = len(A)
    C = [[0]*n for _ in range(n)]
    for i in range(n):
        for k in range(n):
            for j in range(n):
                C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % mod
    return C

메모