====== numtheory.py ======

  * [[numtheory_old|이전 버전 코드]]

===== linear_congruences =====

==== 코드 ====
<dkpr py>
# 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
</dkpr>

==== 설명 ====
  * [[ps:연립 선형 합동식]] 참고
  * 함수 이름에 고민을 많이 했다. 연립 합동식을 정확히 쓰면 "System of linear congruences"나 "Simultaneous linear congruences"가 되는데 둘다 너무 길어서, 차라리 알고리즘 이름대로 "chinese_remainder"를 쓸까 하다가, 결국 지금의 이름으로 결정.
  * 처음 버전에서는 모듈러가 모두 서로소인 경우에 대해서만 해결하도록 작성했다가, 모듈러가 서로소가 아닌 경우까지 확장했다. 그렇게 바뀌면서 리턴값도 가능한 해중의 최솟값만 리턴해주던 것에서 모듈러들의 lcm까지도 함께 리턴해주도록 바뀌었다.


==== 이 코드를 사용하는 문제 ====
---- struct table ----
schema: ps
cols: site, prob_id, %title%, prob_level
filter: teflib ~ *linear_congruences*
csv: 0
----

===== prime_list =====

==== 코드 ====
<dkpr py>
# 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]])
</dkpr>

==== 설명 ====
  * [[ps:소수 목록 구하기]] 참고
  * v1.2x 까지는 2-wheel을 사용했으나, v1.3에서부터 2,3-wheel을 사용
  * 인자를 한개만 주면, [2, a]까지의 소수목록을, 인자 두개를 주면 [a, b]까지의 소수목록을 준다.
  * False의 리스트를 만들어서 slice assignment를 사용하는 것이, 원소 하나씩 대입하는 것보다 빠르다.

==== 이 코드를 사용하는 문제 ====
---- struct table ----
schema: ps
cols: site, prob_id, %title%, prob_level
filter: teflib ~ *prime_list*
csv: 0
----

===== is_prime =====

==== 코드 ====
<dkpr py>
# 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
</dkpr>

==== 설명 ====
  * [[ps:소수 판별#밀러-라빈 알고리즘]] 참고
  * 이 코드를 복사해서 사용하려면, 모듈레벨 상수인 _MILLER_RABIN_THRE, _BASE_FOR_INT32, _BASE_FOR_INT64 도 필요하다.
  

==== 이 코드를 사용하는 문제 ====
---- struct table ----
schema: ps
cols: site, prob_id, %title%, prob_level
filter: teflib ~ *[numtheory.is_prime]*
csv: 0
----



===== prime_factorization =====

==== 코드 ====
<dkpr py>
# 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
</dkpr>

==== 설명 ====
  * [[ps:소인수분해#pollard_s_rho|폴라드 로]] 참고.
  

==== 이 코드를 사용하는 문제 ====
---- struct table ----
schema: ps
cols: site, prob_id, %title%, prob_level
filter: teflib ~ *[numtheory.prime_factorization]*
csv: 0
----


===== all_divisors =====

==== 코드 ====
<dkpr py>
# 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)]
</dkpr>

==== 설명 ====
  * 소인수분해 결과를 이용해서 모든 약수를 생성해서, 리스트에 담아 리턴.
  * 반환된 리스트 안의 약수는 크기순으로 정렬되어 있지 않다.

==== 이 코드를 사용하는 문제 ====
---- struct table ----
schema: ps
cols: site, prob_id, %title%, prob_level
filter: teflib ~ *[numtheory.all_divisors]*
csv: 0
----