#408436 (Rust) No.950 行列累乗

提出ソース

結果

問題	No.950 行列累乗
ユーザー	akakimidori
提出日時	2019-12-14 04:16:15
言語	Rust (1.77.0 + proconio)
結果	AC
実行時間	155 ms / 2,000 ms
コード長	5,881 bytes
コンパイル時間	14,893 ms
コンパイル使用メモリ	379,260 KB
実行使用メモリ	11,116 KB
最終ジャッジ日時	2024-06-28 02:37:01
合計ジャッジ時間	18,528 ms
ジャッジサーバーID （参考情報）	judge5 / judge4

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	28 ms 5,248 KB
testcase_01	AC	27 ms 5,248 KB
testcase_02	AC	1 ms 5,376 KB
testcase_03	AC	1 ms 5,376 KB
testcase_04	AC	1 ms 5,376 KB
testcase_05	AC	36 ms 5,376 KB
testcase_06	AC	1 ms 5,376 KB
testcase_07	AC	1 ms 5,376 KB
testcase_08	AC	1 ms 5,376 KB
testcase_09	AC	1 ms 5,376 KB
testcase_10	AC	1 ms 5,376 KB
testcase_11	AC	1 ms 5,376 KB
testcase_12	AC	1 ms 5,376 KB
testcase_13	AC	1 ms 5,376 KB
testcase_14	AC	1 ms 5,376 KB
testcase_15	AC	1 ms 5,376 KB
testcase_16	AC	1 ms 5,376 KB
testcase_17	AC	1 ms 5,376 KB
testcase_18	AC	1 ms 5,376 KB
testcase_19	AC	1 ms 5,376 KB
testcase_20	AC	1 ms 5,376 KB
testcase_21	AC	39 ms 5,376 KB
testcase_22	AC	55 ms 7,172 KB
testcase_23	AC	37 ms 5,376 KB
testcase_24	AC	46 ms 7,164 KB
testcase_25	AC	30 ms 5,376 KB
testcase_26	AC	43 ms 5,376 KB
testcase_27	AC	28 ms 5,376 KB
testcase_28	AC	50 ms 7,180 KB
testcase_29	AC	45 ms 7,180 KB
testcase_30	AC	46 ms 7,168 KB
testcase_31	AC	42 ms 7,112 KB
testcase_32	AC	65 ms 11,116 KB
testcase_33	AC	48 ms 7,168 KB
testcase_34	AC	64 ms 11,092 KB
testcase_35	AC	48 ms 7,168 KB
testcase_36	AC	86 ms 5,376 KB
testcase_37	AC	115 ms 5,376 KB
testcase_38	AC	120 ms 5,376 KB
testcase_39	AC	91 ms 5,376 KB
testcase_40	AC	51 ms 5,376 KB
testcase_41	AC	50 ms 5,376 KB
testcase_42	AC	40 ms 5,376 KB
testcase_43	AC	68 ms 11,092 KB
testcase_44	AC	67 ms 11,092 KB
testcase_45	AC	67 ms 11,116 KB
testcase_46	AC	69 ms 10,992 KB
testcase_47	AC	37 ms 5,376 KB
testcase_48	AC	67 ms 11,116 KB
testcase_49	AC	67 ms 11,104 KB
testcase_50	AC	68 ms 10,988 KB
testcase_51	AC	68 ms 10,988 KB
testcase_52	AC	153 ms 5,376 KB
testcase_53	AC	155 ms 5,376 KB
testcase_54	AC	1 ms 5,376 KB
testcase_55	AC	1 ms 5,376 KB
testcase_56	AC	40 ms 5,376 KB
testcase_57	AC	1 ms 5,376 KB
testcase_58	AC	1 ms 5,376 KB
testcase_59	AC	145 ms 5,376 KB
testcase_60	AC	151 ms 5,376 KB

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

ソースコード

raw source code

use std::io::Read;

fn mod_pow(r: u64, mut n: u64, p: u64) -> u64 {
    assert!(r < p);
    let mut t = 1;
    let mut s = r;
    while n > 0 {
        if n & 1 == 1 {
            t = t * s % p;
        }
        s = s * s % p;
        n >>= 1;
    }
    t
}

fn inv(a: u64, p: u64) -> u64 {
    assert!(0 < a && a < p);
    mod_pow(a, p - 2, p)
}

fn log(x: u64, a: u64, p: u64) -> u64 {
    assert!(0 < x && x < p && a < p);
    let sq = (p as f64).sqrt() as u64 + 1;
    let mut map = std::collections::HashMap::new();
    for i in (0..sq).rev() {
        map.insert(a * mod_pow(inv(x, p), i, p) % p, i);
    }
    let mut k = 1;
    let x = mod_pow(x, sq, p);
    let mut ans = p;
    for i in 0..sq {
        if let Some(&v) = map.get(&k) {
            let v = (i * sq + v) % (p - 1);
            ans = std::cmp::min(ans, v);
        }
        k = k * x % p;
    }
    ans
}

type M = [[u64; 2]; 2];

fn det(a: &M, p: u64) -> u64 {
    (a[0][0] * a[1][1] % p + p - a[0][1] * a[1][0] % p) % p
}

fn mat_inv(a: &M, p: u64) -> M {
    matmul_s(&[[a[1][1], (p - a[0][1]) % p], [(p - a[1][0]) % p, a[0][0]]], inv(det(a, p), p), p)
}

fn matmul(a: &M, b: &M, p: u64) -> M {
    let mut c = [[0; 2]; 2];
    for (c, a) in c.iter_mut().zip(a.iter()) {
        for (a, b) in a.iter().zip(b.iter()) {
            for (c, b) in c.iter_mut().zip(b.iter()) {
                *c = (*c + *a * *b) % p;
            }
        }
    }
    c
}

fn matmul_s(a: &M, x: u64, p: u64) -> M {
    let mut c = [[0; 2]; 2];
    for (c, a) in c.iter_mut().zip(a.iter()) {
        for (c, a) in c.iter_mut().zip(a.iter()) {
            *c = *a * x % p;
        }
    }
    c
}

fn mat_pow(a: &M, mut n: u64, p: u64) -> M {
    let mut t = [[0; 2]; 2];
    t[0][0] = 1;
    t[1][1] = 1;
    let mut s = a.clone();
    while n > 0 {
        if n & 1 == 1 {
            t = matmul(&t, &s, p);
        }
        s = matmul(&s, &s, p);
        n >>= 1;
    }
    t
}

/*
 * https://twitter.com/maspy_stars/status/1205499459993362432
 * を参考に
 * A^2 = tr(A)A - det(A)I なので
 * det(A) == 0 のときはA^n = (tr(A))^(n - 1)A なので適当に離散対数を解く
 * det(A) != 0 のときはA^n = B から(A^r)(A^qk) = B, det(A^r) = det(B), det(A^q) = 1 
 * となるようなq,rを求める。
 * A をA^q, B をBA^(-r)と置き換えて
 * A^n = B , det(A) = det(B) = 1 が解ければ良い
 * F_p云々はよくわからないがA^n = x * A + y * I とかけることと
 * detA = 1 の条件から x^2 + tr(A)xy + y^2 = 1 (mod p) 
 * という関係式が立ってxに対してyが高々2通りになることから2p程度で抑えられる?
 * */

fn run() {
    let mut s = String::new();
    std::io::stdin().read_to_string(&mut s).unwrap();
    let mut it = s.trim().split_whitespace();
    let p: u64 = it.next().unwrap().parse().unwrap();
    let mut a: M = [[0; 2]; 2];
    let mut b: M = [[0; 2]; 2];
    for i in 0..2 {
        for j in 0..2 {
            a[i][j] = it.next().unwrap().parse().unwrap();
        }
    }
    for i in 0..2 {
        for j in 0..2 {
            b[i][j] = it.next().unwrap().parse().unwrap();
        }
    }
    if a == b {
        println!("1");
        return;
    }
    let det_a = det(&a, p);
    if det_a == 0 {
        let x = (a[0][0] + a[1][1]) % p;
        if x == 0 {
            if b == [[0; 2]; 2] {
                println!("2");
            } else {
                println!("-1");
            }
            return;
        }
        let mut y = p;
        for i in 0..2 {
            for j in 0..2 {
                if a[i][j] != 0 {
                    let m = log(x, b[i][j] * inv(a[i][j], p) % p, p);
                    y = std::cmp::min(y, m);
                }
            }
        }
        if y < p && b == matmul_s(&a, mod_pow(x, y, p), p) {
            println!("{}", y + 1);
        } else {
            println!("-1");
        }
    } else {
        let det_b = det(&b, p);
        if det_b == 0 {
            println!("-1");
            return;
        }
        let v = log(det_a, det_b, p);
        if v >= p {
            println!("-1");
            return;
        }
        let ia = mat_inv(&a, p);
        assert!(matmul(&a, &ia, p) == [[1, 0], [0, 1]]);
        let b = matmul(&b, &mat_pow(&ia, v, p), p);
        assert!(det(&b, p) == 1);
        let mut phi = p - 1;
        let mut factor = vec![];
        for k in 2.. {
            if k * k > phi {
                if phi > 1 {
                    factor.push(phi);
                }
                break;
            }
            if phi % k == 0 {
                factor.push(k);
                while phi % k == 0 {
                    phi /= k;
                }
            }
        }
        let mut phi = p - 1;
        for &f in &factor {
            while phi % f == 0 && mod_pow(det_a, phi / f, p) == 1 {
                phi /= f;
            }
        }
        let phi = phi;
        let a = mat_pow(&a, phi, p);
        assert!(det(&a, p) == 1);
        let ia = mat_inv(&a, p);
        assert!(matmul(&a, &ia, p) == [[1, 0], [0, 1]]);
        let mut map = std::collections::HashMap::new();
        let sq = ((2 * p) as f64).sqrt() as u64 + 100;
        let mut m = matmul(&b, &mat_pow(&ia, sq, p), p);
        for i in (if v == 0 {1} else {0}..(sq + 1)).rev() {
            map.insert(m, i);
            m = matmul(&m, &a, p);
        }
        let mut k = [[1, 0], [0, 1]];
        let x = mat_pow(&a, sq, p);
        let mut ans = sq * sq;
        for i in 0..(sq + 1) {
            if let Some(&y) = map.get(&k) {
                ans = std::cmp::min(ans, sq * i + y);
            }
            k = matmul(&k, &x, p);
        }
        if ans < sq * sq {
            let ans = phi.checked_mul(ans).expect("overflow") + v;
            println!("{}", ans);
        } else {
            println!("-1");
        }
    }
}

fn main() {
    run();
}

yukicoder

結果

テストケース

ソースコード