ps:teflib:numtheory
목차
numtheory.py
linear_congruences
코드
# N linear_congruences
# I {"version": "2.2", "import": ["math"], "abc": ["Iterable"]}
def linear_congruences(a_vals: Iterable[int],
m_vals: Iterable[int],
*,
coprime_moduli: bool = False) -> tuple[int, int]:
"""Returns a solution of the given system of linear congruences.
This function finds the smallest non-negative integer x0 and m, where any
x_k = x0 + k*m satisfies following simultaneous linear congruence equations:
x % m_0 = a_0
x % m_1 = a_1
...
x % m_n = a_n
The argument 'coprime_moduli' should be set True if all m_i's are pairwise
coprime. It runs faster by using Chinese Remainder Theorem.
"""
if coprime_moduli:
m_list = list(m_vals)
m = 1
for m_i in m_list:
m *= m_i
a = sum((p := m // m_i) * pow(p, -1, m_i) * a_i
for a_i, m_i in zip(a_vals, m_list)) % m
else:
a, m = 0, 1
for b, n in zip(a_vals, m_vals):
g = math.gcd(m, n)
l, mod = divmod(a - b, g)
if mod:
raise ValueError('No solution')
um = pow(m // g, -1, n // g) * m
m *= n // g
a = (a - l * um) % m
return a, m
설명
- 연립 선형 합동식 참고
- 함수 이름에 고민을 많이 했다. 연립 합동식을 정확히 쓰면 “System of linear congruences”나 “Simultaneous linear congruences”가 되는데 둘다 너무 길어서, 차라리 알고리즘 이름대로 “chinese_remainder”를 쓸까 하다가, 결국 지금의 이름으로 결정.
- 처음 버전에서는 모듈러가 모두 서로소인 경우에 대해서만 해결하도록 작성했다가, 모듈러가 서로소가 아닌 경우까지 확장했다. 그렇게 바뀌면서 리턴값도 가능한 해중의 최솟값만 리턴해주던 것에서 모듈러들의 lcm까지도 함께 리턴해주도록 바뀌었다.
이 코드를 사용하는 문제
prime_list
코드
# N prime_list
# I {"version": "1.3", "import": ["math"]}
def prime_list(a: int, b: int = -1) -> list[int]:
"""Returns a list of prime numbers in the given range, [1, a] or [a, b]."""
beg, end = (1, a + 1) if b < 0 else (min(a, b), max(a, b) + 1)
if end < 5:
return [i for i in range(beg, end) if i in (2, 3)]
n = end + 6 - end % 6
sieve = [False] + [True] * (n // 3 - 1)
for i in range(math.isqrt(n) // 3 + 1):
if sieve[i]:
d, s, j = (k := 1 | 3 * i + 1) * 2, k * k, k * (k + 4 - 2 * (i & 1))
sieve[s // 3::d] = [False] * ((n // 6 - s // 6 - 1) // k + 1)
sieve[j // 3::d] = [False] * ((n // 6 - j // 6 - 1) // k + 1)
b, e = (beg | 1) // 3, n // 3 - 2 + (end % 6 > 1)
return ([p for p in (2, 3) if p >= beg] +
[1 | 3 * i + 1 for i in range(b, e) if sieve[i]])
설명
- 소수 목록 구하기 참고
- v1.2x 까지는 2-wheel을 사용했으나, v1.3에서부터 2,3-wheel을 사용
- 인자를 한개만 주면, [2, a]까지의 소수목록을, 인자 두개를 주면 [a, b]까지의 소수목록을 준다.
- False의 리스트를 만들어서 slice assignment를 사용하는 것이, 원소 하나씩 대입하는 것보다 빠르다.
이 코드를 사용하는 문제
| 출처 | 문제 번호 | Page | 레벨 |
|---|---|---|---|
| BOJ | 1644 | 소수의 연속합 | 골드 3 |
| BOJ | 23633 | 소수 징글벨 | 골드 3 |
| BOJ | 1153 | 네 개의 소수 | 골드 4 |
| BOJ | 4913 | 페르마의 크리스마스 정리 | 골드 4 |
| BOJ | 17104 | 골드바흐 파티션 2 | 다이아몬드 5 |
| BOJ | 17103 | 골드바흐 파티션 | 실버 2 |
| BOJ | 6219 | 소수의 자격 | 실버 3 |
| BOJ | 31216 | 슈퍼 소수 | 실버 5 |
| BOJ | 26973 | Circular Barn | 플래티넘 3 |
| BOJ | 11439 | 이항 계수 5 | 플래티넘 4 |
is_prime
코드
# N is_prime
# I {"version": "1.3"}
@functools.cache
def is_prime(n: int) -> bool:
"""Check if n is prime number, in O(logn).
This uses Miller-Rabin algorithm.
The answer is correct for n<2^64. For n>2^64, the answer could be incorrect
in a low probability (p <= 1/(4^7)).
"""
if n <= _MILLER_RABIN_THRE:
if n <= 2 or n % 2 == 0:
return n == 2
return all(n % i for i in range(3, math.isqrt(n) + 1, 2))
bases = _BASE_FOR_INT32 if n.bit_length() < 32 else _BASE_FOR_INT64
s = ((n - 1) & (1 - n)).bit_length() - 1
d = n >> s
for a in bases:
if (num := pow(a, d, n)) not in (1, n - 1) and all(
(num := num * num % n) != n - 1 for _ in range(s - 1)
):
return False
return True
설명
- 밀러-라빈 알고리즘 참고
- 이 코드를 복사해서 사용하려면, 모듈레벨 상수인 _MILLER_RABIN_THRE, _BASE_FOR_INT32, _BASE_FOR_INT64 도 필요하다.
이 코드를 사용하는 문제
prime_factorization
코드
# N prime_factorization
# I {"version": "1.12", "deps": ["is_prime"]}
def _find_factor(n: int) -> int:
"""Return a random factor of n, using Brent-Pollard Rho method.
n should not be prime."""
m = int(n**0.125) + 1
for c in itertools.count(start=1):
x, r, g, q, xs, y = 2, 1, 1, 1, 0, 0
while g == 1:
y = x
for _ in range(r >> 1, (3 * r) >> 2):
x = (x * x + c) % n
for k in range((3 * r) >> 2, r, m):
xs = x
for _ in range(min(m, r - k)):
x = (x * x + c) % n
q = q * abs(y - x) % n
if (g := math.gcd(q, n)) != 1:
break
r += r
if g == n:
xs = (xs * xs + c) % n
while (g := math.gcd(xs - y, n)) == 1:
xs = (xs * xs + c) % n
if g != n:
return g
return -1
def prime_factorization(n: int) -> dict[int, int]:
"""Return {p1:e1, ..., pk:ek}, where n = p1^e1 * ... * pk^ek, in O(n^1/4).
For n == 1, return an empty dict."""
assert n > 0, f'Should be n>0. ({n=})'
factorization = collections.defaultdict(int)
if n % 2 == 0:
power_two = factorization[2] = (n & -n).bit_length() - 1
n >>= power_two
if n == 1:
return factorization
factors = [n]
while factors:
n = factors.pop()
while not is_prime(n):
g = _find_factor(n)
factors.append(n // g)
n = g
factorization[n] += 1
return factorization
설명
- 폴라드 로 참고.
이 코드를 사용하는 문제
| 출처 | 문제 번호 | Page | 레벨 |
|---|---|---|---|
| BOJ | 9326 | MI6 | 골드 3 |
| BOJ | 17633 | 제곱수의 합 (More Huge) | 다이아몬드 4 |
| BOJ | 8187 | Divine Divisor | 다이아몬드 5 |
| BOJ | 17646 | 제곱수의 합 2 (More Huge) | 루비 4 |
all_divisors
코드
# N all_divisors
# I {"version": "1.0", "lastupdated": "20260125", "deps": ["prime_factorization"]}
def all_divisors(n):
"""Return a list of all divisors of n, in O(n^1/3)."""
prime_powers = []
for p, e in prime_factorization(n).items():
prime_powers.append([x := 1] + [x := x * p for _ in range(e)])
return [math.prod(t) for t in itertools.product(*prime_powers)]
설명
- 소인수분해 결과를 이용해서 모든 약수를 생성해서, 리스트에 담아 리턴.
- 반환된 리스트 안의 약수는 크기순으로 정렬되어 있지 않다.
이 코드를 사용하는 문제
| 출처 | 문제 번호 | Page | 레벨 |
|---|---|---|---|
| BOJ | 4464 | Pride and Prejudice and Zombies | 다이아몬드 5 |
| BOJ | 23362 | Rasterized Lines | 플래티넘 1 |
ps/teflib/numtheory.txt · 마지막으로 수정됨: 2026/02/01 13:50 저자 teferi

토론