#681042 (PyPy3) No.981 一般冪乗根

提出ソース

結果

問題	No.981 一般冪乗根
ユーザー	hari64
提出日時	2021-07-22 18:09:46
言語	PyPy3 (7.3.15)
結果	AC
実行時間	3,688 ms / 6,000 ms
コード長	6,784 bytes
コンパイル時間	175 ms
コンパイル使用メモリ	82,504 KB
実行使用メモリ	216,004 KB
最終ジャッジ日時	2024-04-18 14:41:56
合計ジャッジ時間	158,693 ms
ジャッジサーバーID （参考情報）	judge2 / judge5

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	3,323 ms 200,024 KB
testcase_01	AC	3,309 ms 203,856 KB
testcase_02	AC	3,288 ms 199,120 KB
testcase_03	AC	3,280 ms 198,508 KB
testcase_04	AC	3,310 ms 197,084 KB
testcase_05	AC	199 ms 77,484 KB
testcase_06	AC	205 ms 77,240 KB
testcase_07	AC	201 ms 77,636 KB
testcase_08	AC	199 ms 77,780 KB
testcase_09	AC	201 ms 77,340 KB
testcase_10	AC	216 ms 77,644 KB
testcase_11	AC	202 ms 77,584 KB
testcase_12	AC	207 ms 77,820 KB
testcase_13	AC	196 ms 77,560 KB
testcase_14	AC	198 ms 77,620 KB
testcase_15	AC	204 ms 77,292 KB
testcase_16	AC	199 ms 77,408 KB
testcase_17	AC	201 ms 77,432 KB
testcase_18	AC	200 ms 77,436 KB
testcase_19	AC	192 ms 77,616 KB
testcase_20	AC	199 ms 77,688 KB
testcase_21	AC	195 ms 77,192 KB
testcase_22	AC	200 ms 77,724 KB
testcase_23	AC	197 ms 77,400 KB
testcase_24	AC	196 ms 77,740 KB
testcase_25	AC	3,235 ms 202,120 KB
testcase_26	AC	3,688 ms 216,004 KB
testcase_27	AC	109 ms 77,604 KB
testcase_28	AC	3,087 ms 194,692 KB
evil_60bit1.txt	TLE	-
evil_60bit2.txt	MLE	-
evil_60bit3.txt	MLE	-
evil_hack	AC	154 ms 62,396 KB
evil_hard_random	MLE	-
evil_hard_safeprime.txt	MLE	-
evil_hard_tonelli0	MLE	-
evil_hard_tonelli1	MLE	-
evil_hard_tonelli2	MLE	-
evil_hard_tonelli3	MLE	-
evil_sefeprime1.txt	MLE	-
evil_sefeprime2.txt	MLE	-
evil_sefeprime3.txt	MLE	-
evil_tonelli1.txt	MLE	-
evil_tonelli2.txt	MLE	-

権限があれば一括ダウンロードができます

ソースコード

raw source code

def is_prime(n: int) -> bool: #ミラー–ラビン素数判定法
    # https://qiita.com/srtk86/items/609737d50c9ef5f5dc59
    # https://en.wikipedia.org/wiki/Strong_pseudoprime
    assert isinstance(n, int) and 0 < n
    if n == 2:
        return True
    if n == 1 or n & 1 == 0:
        return False
    d = (n - 1) >> 1
    while d & 1 == 0:
        d >>= 1  # n-1=(2**s)*d (dは奇数)
    L=[a for a in (2, 325, 9375, 28178,450775, 9780504, 1795265022) if a<n]
    if n>=2**64:  # これはwikiによる
        from random import randint
        L += [randint(1, n-1) for _ in range(100)]
    for a in L:  # nが素数ならばa^d≡1 (mod p) もしくは
        t = d  # a^((2^r)*d)≡-1 (mod p) が成立すべき
        y = pow(a, t, n)
        while t != n - 1 and y != 1 and y != n - 1:
            y = (y * y) % n
            t <<= 1
        if y != n - 1 and t & 1 == 0:
            return False
    else:
        return True


def findFactorRho(n):
    from math import gcd
    m = 1 << n.bit_length() // 8
    for c in range(1, 99):
        def f(x):
            return (x * x + c) % n
        y, r, q, g = 2, 1, 1, 1
        while g == 1:
            x = y
            for i in range(r):
                y = f(y)
            k = 0
            while k < r and g == 1:
                ys = y
                for i in range(min(m, r - k)):
                    y = f(y)
                    q = q * abs(x - y) % n
                g = gcd(q, n)
                k += m
            r <<= 1
        if g == n:
            g = 1
            while g == 1:
                ys = f(ys)
                g = gcd(abs(x - ys), n)
        if g < n:
            if is_prime(g):
                return g
            elif is_prime(n // g):
                return n // g
            return findFactorRho(g)


def prime_factors(n, return_e: bool = False):
    i = 2
    ret = {}
    rhoFlg = 0
    while i*i <= n:
        k = 0
        while n % i == 0:
            n //= i
            k += 1
        if k:
            ret[i] = k
        i += 1 + i % 2
        if i == 101 and n >= 2 ** 20:
            while n > 1:
                if is_prime(n):
                    ret[n], n = 1, 1
                else:
                    rhoFlg = 1
                    j = findFactorRho(n)
                    k = 0
                    while n % j == 0:
                        n //= j
                        k += 1
                    ret[j] = k
    if n > 1:
        ret[n] = 1
    if rhoFlg:
        ret = {x: ret[x] for x in sorted(ret)}
    return list(ret.keys())


def primitive_root(p: int) -> int:
    """素数pに対する何らかの原始根rを一つ求める (rの位数はp-1)
    https://manabitimes.jp/math/842
    https://37zigen.com/primitive-root/ 最速ではない
    """
    if p == 2:
        return 1
    from random import randint
    p_1_pfs = prime_factors(p-1, return_e=False)
    cnt = 0
    while cnt < 2*p:
        cnt += 1
        cfr = randint(2, p-1)  # candidates for r
        for pf in p_1_pfs:
            if pow(cfr, (p-1)//pf, p) == 1:  # p-1//pfが位数候補
                break
        else:  # 全ての位数候補に対して、その否定がなされた。
            return cfr
    else:
        raise AssertionError


def discrete_logarithm(x, y, mod):
    """離散対数問題 x^k≡yとなるようなkを求める O(√mod)"""
    babys = {}
    sqrt = int(mod**0.5)+1
    # baby-step
    x_baby = 1
    for i in range(sqrt):
        babys[x_baby] = i  # x^(sqrt-1)まで求めておく
        if x_baby == y:
            return i
        x_baby = x_baby * x % mod
    # giant-step
    x_giant = pow(x_baby, mod-2, mod)  # x^(-sqrt)
    for i in range(1, sqrt+1):
        y = y * x_giant % mod
        if y in babys:  # ここで√modずつ見ていく為に速い
            return babys[y] + i*sqrt
    return -1  # 原始根でない限り、必ず存在するとは言えない


def _inv_gcd(a, b) -> tuple:
    """
    返り値は(gcd(a,b),x) (ただしxa≡g (mod b) 0<=x<b//g)
    計算量 O(log(max(a,b))) ラメの定理などからも分かる通り計算量が非常に小さい
    (0<=a<=b∊N len(str(a))=dの時、gcd算出の為の計算回数は5d以下)
    逆元を求めるならpow(g,-1,b)でもいいが、versionの問題とgとbが互いに素という制約はある
    また、このコードは拡張ユークリッドの互除法の行列表現からも理解可能 再帰でも書ける
    https://github.com/atcoder/ac-library/blob/master/atcoder/internal_math.hpp
    """
    # assert 0 <= a and 1 <= b
    a %= b
    if a == 0:
        return (b, 0)
    s = b  # m0*a≡s (mod b) これらの性質をm0,m1は満たしていることに注意
    t = a  # m1*a≡t (mod b)
    m0 = 0  # s*|m1|+t*|m0|<=b
    m1 = 1  # また、a<bよりt<sが成立
    while t:
        # この三行は互除法そのもの
        u = s // t
        s -= t * u
        s, t = t, s
        # m0とm1が先ほど挙げた性質を保持し続ける様に変更する oがoldm、nがnewを示すとして
        # o_m0*a=o_s o_m1*a=o_t (冒頭にあげた性質)
        # n_s=o_t n_t=o_s-o_t*u (互除法による変更)
        # o_m1*a=n_s (o_m0-o_m1*u)*a=n_t (代入して整理)
        # これを先の性質の式と見比べるとn_m0=o_m1 n_m1=o_m0-o_m1*uを得る これが下の式の正体
        # 三つ目の性質も保たれていることは代入すればすぐわかる
        m0 -= m1 * u
        m0, m1 = m1, m0
    if m0 < 0:  # 三つ目の性質を利用するとu*|n_m0|<=b//s=b//g⇒|n_mo|<b/g (∵u>=0)
        m0 += b // s
    return (s, m0)


def eea(a: int, b: int) -> tuple:
    """
    ax+by=gcd(a,b)なる(x,y)を求める(拡張されたユークリッドの互除法)
    (extended_euclidean_algorithm)(元のユークリッドの互除法はgcdを求める手法を指す)
    https://qiita.com/drken/items/b97ff231e43bce50199a
    https://ja.wikipedia.org/wiki/ユークリッドの互除法 <-英語版が優秀すぎ
    https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm
    また、以上の話題は連分数展開とも関連があるらしい
    """
    g, x = _inv_gcd(a, abs(b))
    y = (b//abs(b))*(g-a*x)//abs(b)
    assert (g-a*x) % abs(b) == 0
    return (x, y)


def main():
    import sys
    from math import gcd
    input = sys.stdin.buffer.readline
    T = int(input())
    for _ in range(T):
        p, k, a = map(int, input().split())
        g = primitive_root(p)
        y = discrete_logarithm(g, a, p)
        gkp = gcd(k, p-1)
        if y % gkp != 0:
            print(-1)
            continue
        z = (y//gkp)*eea(k//gkp, -(p-1)//gkp)[0]
        x = pow(g, z % (p-1), p)
        print(x)


if __name__ == '__main__':
    main()

yukicoder

結果

テストケース

ソースコード