__init__.py

import logging
from math import gcd
from math import isqrt


def int_to_bits_le(i, count):
    """
    Converts an integer to bits, little endian.
    :param i: the integer
    :param count: the number of bits
    :return: the bits
    """
    bits = []
    for _ in range(count):
        bits.append(i & 1)
        i >>= 1

    return bits


def bits_to_int_le(bits, count):
    """
    Converts bits to an integer, little endian
    :param bits: the bits
    :param count: the number of bits
    :return: the integer
    """
    i = 0
    for k in range(count):
        i |= (bits[k] & 1) << k

    return i


def floor_div(a, b):
    """
    Returns floor(a / b), works with large integers.
    :param a: a
    :param b: b
    :return: floor(a / b)
    """
    return a // b


def ceil_div(a, b):
    """
    Returns ceil(a / b), works with large integers.
    :param a: a
    :param b: b
    :return: ceil(a / b)
    """
    return a // b + (a % b > 0)


def is_square(x):
    """
    Returns the square root of x if x is a perfect square, or None otherwise.
    :param x: x
    :return: the square root of x or None
    """
    y = isqrt(x)
    return y if y ** 2 == x else None


def symmetric_mod(x, m):
    """
    Computes the symmetric modular reduction.
    :param x: the number to reduce
    :param m: the modulus
    :return: x reduced in the interval [-m/2, m/2]
    """
    return int((x + m + m // 2) % m) - int(m // 2)


def solve_congruence(a, b, m):
    """
    Solves a congruence of the form ax = b mod m.
    :param a: the parameter a
    :param b: the parameter b
    :param m: the modulus m
    :return: a generator generating solutions for x
    """
    g = gcd(a, m)
    a //= g
    b //= g
    n = m // g
    for i in range(g):
        yield (pow(a, -1, n) * b + i * n) % m


def divisors(factors):
    """
    Computes all divisors from a list of factors
    :param factors: the factors (tuples of primes and exponents)
    :return: a generator generating divisors
    """
    divisors = [1]
    yield 1
    for p, e in factors:
        new = []
        for d in divisors:
            for k in range(1, e + 1):
                d_ = p ** k * d
                new.append(d_)
                yield int(d_)

        divisors += new


def make_square_free(x, factors):
    """
    For any integer x, removes all square factors.
    :param x: the value x
    :param factors: the factors of x
    :return: a square-free integer y, corresponding to x with all square factors removed
    """
    for p, e in factors:
        while e > 0 and e % 2 == 0:
            e -= 2
            x //= p
            x //= p
    return int(x)


def roots_of_unity(ring, l, r):
    """
    Generates r-th roots of unity in a ring, with r | l.
    :param ring: the ring, with order n
    :param l: the Carmichael lambda of n
    :param r: r
    :return: a generator generating the roots of unity
    """
    assert l % r == 0, "r should divide l"

    x = ring(2)
    while (g := x ** (l // r)) == 1:
        x += 1

    for i in range(r):
        yield int(g ** i)


def rth_roots(Fq, delta, r):
    """
    Uses the Adleman-Manders-Miller algorithm to extract r-th roots in Fq, with r | q - 1.
    More information: Cao Z. et al., "Adleman-Manders-Miller Root Extraction Method Revisited" (Table 4)
    :param Fq: the field Fq
    :param delta: the r-th residue delta
    :param r: the r
    :return: a generator generating the rth roots
    """
    delta = Fq(delta)
    q = Fq.order()
    assert (q - 1) % r == 0, "r should divide q - 1"

    p = Fq(1)
    while p ** ((q - 1) // r) == 1:
        p = Fq.random_element()

    t = 0
    s = q - 1
    while s % r == 0:
        t += 1
        s //= r

    k = 1
    while (k * s + 1) % r != 0:
        k += 1
    alpha = (k * s + 1) // r

    a = p ** (pow(r, t - 1, q - 1) * s)
    b = delta ** (r * alpha - 1)
    c = p ** s
    h = 1
    for i in range(1, t):
        d = b ** pow(r, t - 1 - i, q - 1)
        logging.debug(f"Computing the discrete logarithm for {i = }, this may take a long time...")
        j = 0 if d == 1 else -d.log(a)
        b *= (c ** r) ** j
        h *= c ** j
        c **= r

    root = int(delta ** alpha * h)
    for primitive_root in roots_of_unity(Fq, q - 1, r):
        yield root * primitive_root % q


def modinv_range(n, p):
    """
    Computes the modular inverses of the numbers in the range (1, n] (exclusive), mod p.
    More information: grhkm, "[Tutorial] Calculate modulo inverses efficiently!" (Codeforces)
    :param n: the n
    :param p: the modulus
    :return: a generator generating the modular inverses of 1, 2... n - 1 mod p
    """
    inv = [0] * n
    inv[1] = 1
    yield inv[1]
    for i in range(2, n):
        inv[i] = (p - p // i) * inv[p % i] % p
        yield inv[i]


def modinv(a, p):
    """
    Computes the modular inverses a list of numbers mod p.
    More information: grhkm, "[Tutorial] Calculate modulo inverses efficiently!" (Codeforces)
    :param a: the list of numbers
    :param p: the modulus
    :return: a generator generating the modular inverses of a1, a2... mod p
    """
    n = len(a)
    pre = [0] * n
    pre[0] = 1
    suf = [0] * n
    suf[n - 1] = 1
    prod = 1
    for i in range(n - 1):
        pre[i + 1] = pre[i] * a[i] % p
        suf[n - i - 2] = suf[n - i - 1] * a[n - i - 1] % p
        prod = prod * a[i] % p

    prod = prod * a[n - 1] % p
    prod = pow(prod, -1, p)
    for i in range(n):
        yield (pre[i] * suf[i] % p) * prod % p