結果

問題 No.981 一般冪乗根
ユーザー akakimidori
提出日時 2025-12-08 23:27:43
言語 Rust
(1.83.0 + proconio)
結果
RE  
実行時間 -
コード長 4,143 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 21,182 ms
コンパイル使用メモリ 402,020 KB
実行使用メモリ 7,852 KB
最終ジャッジ日時 2025-12-08 23:28:41
合計ジャッジ時間 58,107 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 3 RE * 41
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: unused import: `std::io::Write`
 --> src/main.rs:2:5
  |
2 | use std::io::Write;
  |     ^^^^^^^^^^^^^^
  |
  = note: `#[warn(unused_imports)]` on by default

warning: type alias `Map` is never used
 --> src/main.rs:4:6
  |
4 | type Map<K, V> = BTreeMap<K, V>;
  |      ^^^
  |
  = note: `#[warn(dead_code)]` on by default

warning: type alias `Set` is never used
 --> src/main.rs:5:6
  |
5 | type Set<T> = BTreeSet<T>;
  |      ^^^

warning: type alias `Deque` is never used
 --> src/main.rs:6:6
  |
6 | type Deque<T> = VecDeque<T>;
  |      ^^^^^

ソースコード

diff #
raw source code 

use std::collections::*;
use std::io::Write;

type Map<K, V> = BTreeMap<K, V>;
type Set<T> = BTreeSet<T>;
type Deque<T> = VecDeque<T>;

fn main() {
    input! {
        t: usize,
        ask: [(i64, i64, i64); t],
    }
    for (p, k, a) in ask {
        assert!(k <= 1000);
        if let Some(x) = kth_root(a, k, p) {
            assert!(pow_mod(x, k, p) == a);
            println!("{}", x);
        } else {
            println!("-1");
        }
    }
}

// x^k = a (mod p)
// となるようなxが存在するなら一つ返す
// p: 素数
pub fn kth_root(a: i64, k: i64, p: i64) -> Option<i64> {
    if a < 0 || a >= p {
        return None;
    }
    assert!(p >= 2 && k > 0);
    if a == 0 {
        return Some(a);
    }
    let (x, _, g) = ext_gcd(k, p - 1);
    if pow_mod(a, (p - 1) / g, p) != 1 {
        return None;
    }
    let x = x.rem_euclid((p - 1) / g);
    let a = pow_mod(a, x, p);
    let k = g;
    let f = factor(ext_gcd(k, (p - 1) / k).2);
    let mut ans = {
        let mut q = p - 1;
        loop {
            let g = ext_gcd(q, k).2;
            if g == 1 {
                let (x, _, _) = ext_gcd(k, q);
                q = x.rem_euclid(q);
                break;
            }
            q /= g;
        }
        pow_mod(a, q, p)
    };
    for &f in f.iter() {
        let mut q = p - 1;
        let mut r = 0;
        while q % f == 0 {
            q /= f;
            r += 1;
        }
        let mut w = (1..).find(|i| pow_mod(*i, (p - 1) / f, p) != 1).unwrap();
        w = pow_mod(w, q, p);
        for i in 1..r {
            let a = pow_mod(a, q * f.pow(r - i - 1), p);
            while a != pow_mod(ans, q * f.pow(r - i), p) {
                ans = mul(ans, w, p);
            }
            w = mul(w, w, p);
        }
    }
    Some(ans)
}

pub const fn ext_gcd(a: i64, b: i64) -> (i64, i64, i64) {
    if b == 0 {
        (a.signum(), 0, a.abs())
    } else {
        let q = a.div_euclid(b);
        let r = a.rem_euclid(b);
        let (x, y, g) = ext_gcd(b, r);
        (y, x - q * y, g)
    }
}

pub fn factor(mut n: i64) -> Vec<i64> {
    let mut f = vec![];
    let mut p = 2;
    while p * p <= n {
        if n % p == 0 {
            f.push(p);
        }
        while n % p == 0 {
            n /= p;
        }
        p += 1;
    }
    if n > 1 {
        f.push(n);
    }
    f
}

pub const fn mul(a: i64, b: i64, m: i64) -> i64 {
    (a as i128 * b as i128 % m as i128) as i64
}

pub const fn pow_mod(mut r: i64, mut n: i64, p: i64) -> i64 {
    r %= p;
    let mut t = 1 % p;
    while n > 0 {
        if n & 1 == 1 {
            t = mul(t, r, p);
        }
        r = mul(r, r, p);
        n >>= 1;
    }
    t
}

// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
#[macro_export]
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

#[macro_export]
macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};
    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

#[macro_export]
macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };
    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };
    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };
    ($iter:expr, bytes) => {
        read_value!($iter, String).bytes().collect::<Vec<u8>>()
    };
    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };
    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}
// ---------- end input macro ----------
0