Codeforces Round #536 (Div. 2) F. Lunar New Year and a Recursive Sequence

暗号とか勉強したことがあれば実家。
https://codeforces.com/contest/1106/problem/F

問題

p = 998244353とする。正の整数列(f_i)_iが以下のように定められている:

f_1 = ... = f_{k-1} = 1
f_k = x
i > kのときf_i = (f_{i-1}^{b_1} f_{i-2}^{b_2} ... f_{i-k}^{b_k}) mod p

f_n = mがわかっている。このとき、xとしてあり得る1以上p未満の値を一つ出力するか、またはそのような値がないことを指摘せよ。

制約

1 <= k <= 100
1 <= b_i < p
k < n <= 10^9
1 <= m < p

解法

f_i = x^{h_i} mod pの形で表すことにすると、h_iは線形漸化式

h_1 = ... = h_{k - 1} = 0
h_k = 1
i > kのときh_i = b_1 * h_{i-1} + b_2 * h_{i-2} + ... + b_k * h_{i-k}

を満たす。このような数列の第n項は行列累乗を行うことでO(k^3 log n)時間で計算できる。
e = h_nと置くことにする。x^n = mなので、これを満たすxを求めたい。
ここで、mod pで見たときに、0以外のすべての整数は単一の整数の冪乗で表現できることを利用する。p = 998244353の場合は3がその性質を満たす、つまり任意の1からp-1までの整数mに対してm = 3^y mod pとなる整数yが存在する。
問題のmについてそのようなyが求められるのであれば、求めたいものはx^n = 3^yとなる整数xである。まずそのようなyを求めよう。Baby-step giant-stepと呼ばれるアルゴリズムを適用する。S = 40000 (~= sqrt(p))として、1, 3^S, 3^{2S}, ..., 3^{(S-1)S}をキー、0, S, 2S, ..., (S-1)Sを値にして、連想配列を作る。mが与えられた時、m, 3m, 3^2 m, ..., 3^{S - 1}mのどれかは3^Sの冪乗の形になる、つまり連想配列の中に含まれる。これの計算量はO(sqrt(p))である。
ここで、3^{p-1} = 1 (mod p)、また3^z = 1 (mod p)ならばzはp-1の倍数であることに気をつけると、上の問題は、ある整数u,vが存在してun = y + (p-1)vであるかどうかを調べることと同値である。これは拡張ユークリッドの互除法でできる。ここで求まったuについてx = 3^u mod pとすればそれが答え。uが複数あり得る場合はどれでもよい。

以上をまとめると、計算量はO(k^3 log n + sqrt(m))であり、十分間に合う。

登場する典型

行列累乗
(Z/pZ)^*が巡回群であること
- あるg (原始根)が存在して、全ての整数はg^k mod pの形で表せる
離散対数問題
- Baby-step giant-stepアルゴリズム

実装上の注意点

uが負の数の場合はp-1の倍数を足して非負にする(1敗)

提出: Submission #49502575 - Codeforces (Rust)

// ModInt 省略

fn squmul(a: &[Vec<i64>], b: &[Vec<i64>], mo: i64) -> Vec<Vec<i64>> {
    let n = a.len();
    let mut ret = vec![vec![0; n]; n];
    for i in 0..n {
        for j in 0..n {
            for k in 0..n {
                ret[i][k] += a[i][j] * b[j][k];
                ret[i][k] %= mo;
            }
        }
    }
    ret
}

fn squpow(a: &[Vec<i64>], mut e: i64, mo: i64) -> Vec<Vec<i64>> {
    let n = a.len();
    let mut sum = vec![vec![0; n]; n];
    for i in 0..n { sum[i][i] = 1; }
    let mut cur = a.to_vec();
    while e > 0 {
        if e % 2 == 1 {
            sum = squmul(&sum, &cur, mo);
        }
        cur = squmul(&cur, &cur, mo);
        e /= 2;
    }
    sum
}

fn extgcd(x: i64, y: i64) -> (i64, i64, i64) {
    if y == 0 {
        return (x, 1, 0);
    }
    let r = x % y;
    let q = x / y;
    let (g, a, b) = extgcd(y, r);
    (g, b, a - b * q)
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        k: usize,
        b: [i64; k],
        n: i64,
        m: i64,
    }
    let mut trans = vec![vec![0; k]; k];
    for i in 0..k - 1 {
        trans[i + 1][i] = 1;
    }
    for j in 0..k {
        trans[k - 1 - j][k - 1] = b[j] % (MOD - 1);
    }
    let exp = squpow(&trans, n - 1, MOD - 1)[k - 1][0];
    // Find x s.t. x^exp = m
    let gen = ModInt::new(3);
    const SQRT: i64 = 40000;
    let mut cur = ModInt::new(1);
    let prog = gen.pow(SQRT);
    let mut tbl = HashMap::new();
    for i in 0..SQRT {
        tbl.insert(cur, SQRT * i);
        cur *= prog;
    }
    let discrete_log = |x: ModInt| {
        for i in 0..SQRT {
            let key = x * gen.pow(i);
            if let Some(&y) = tbl.get(&key) {
                return (y - i + MOD - 1) % (MOD - 1);
            }
        }
        unreachable!();
    };
    let logm = discrete_log(ModInt::new(m));
    let (g, inv, _) = extgcd(exp, MOD - 1);
    if logm % g != 0 {
        puts!("-1\n");
        return;
    }
    let inv = (inv + MOD - 1) % (MOD - 1);
    let ans = (logm / g * inv) % (MOD - 1);
    puts!("{}\n", gen.pow(ans));
}

まとめ

原始根もいずれは誰でも知っている定跡になりそう。