Submission #73110350


Source Code Expand

use proconio::input;

const MOD: i64 = 998244353;

#[derive(Clone, Copy, Debug)]
struct Gauss {
    r: i64, // 実部 mod C
    i: i64, // 虚部 mod C
}

impl Gauss {
    fn new(r: i64, i: i64, c: i64) -> Self {
        Self {
            r: ((r % c) + c) % c,
            i: ((i % c) + c) % c,
        }
    }

    fn mul(self, other: Self, c: i64) -> Self {
        let r = (self.r * other.r - self.i * other.i) % c;
        let i = (self.r * other.i + self.i * other.r) % c;
        Self::new(r, i, c)
    }

    fn pow(mut self, mut exp: i64, c: i64) -> Self {
        let mut res = Self::new(1, 0, c);
        while exp > 0 {
            if exp & 1 == 1 {
                res = res.mul(self, c);
            }
            self = self.mul(self, c);
            exp >>= 1;
        }
        res
    }
}

// z の冪乗列を計算し、周期の開始位置を返す
fn get_path(z: Gauss, c: i64) -> (Vec<Gauss>, usize) {
    let mut path = Vec::new();
    let mut visited = vec![-1i64; (c * c) as usize];
    let mut cur = Gauss::new(1, 0, c);
    let mut step = 0;
    loop {
        let idx = (cur.r * c + cur.i) as usize;
        if visited[idx] != -1 {
            return (path, visited[idx] as usize);
        }
        visited[idx] = step;
        path.push(cur);
        cur = cur.mul(z, c);
        step += 1;
    }
}

fn main() {
    input! {
        a: i64,
        b: i64,
        c: i64,
        m: usize,
        pe: [(i64, i64); m],
    }

    // DP[state]: state = r * C + i に対応する (r, i) mod C の解の個数
    let mut dp = vec![0i64; (c * c) as usize];
    dp[(1 % c * c) as usize] = 1; // 1 = 1^2 + 0^2

    for &(p, e) in &pe {
        let mut ways = Vec::new(); // (state, count)

        if p % 4 == 3 {
            // p ≡ 3 (mod 4): Z[i] で素元のまま
            if e % 2 == 1 {
                println!("0");
                return;
            }
            // p^e = (p^(e/2))^2 + 0^2
            let z = Gauss::new(p % c, 0, c).pow(e / 2, c);
            ways.push((z.r * c + z.i, 1));
        } else if p == 2 {
            // p = 2 = (1+i)(1-i) = -i(1+i)^2
            let z = Gauss::new(1, 1, c).pow(e, c);
            ways.push((z.r * c + z.i, 1));
        } else {
            // p ≡ 1 (mod 4): p = a^2 + b^2 と分解可能
            let (mut a, mut b) = (1i64, 0i64);
            while a * a <= p {
                let rem = p - a * a;
                let s = (rem as f64).sqrt() as i64;
                if s * s == rem {
                    b = s;
                    break;
                }
                a += 1;
            }

            let z1 = Gauss::new(a % c, b % c, c);
            let z2 = Gauss::new(a % c, (-b % c + c) % c, c);
            let (path1, cycle_start1) = get_path(z1, c);
            let (path2, cycle_start2) = get_path(z2, c);
            let cycle_len1 = path1.len() - cycle_start1;
            let cycle_len2 = path2.len() - cycle_start2;

            let mut cnt = vec![0i64; (c * c) as usize];

            // k: (a+bi)^k の指数, e-k: (a-bi)^(e-k) の指数
            let mut add = |k: i64, count: i64| {
                // (a+bi)^k を計算(周期性を利用)
                let idx1 = if k < path1.len() as i64 {
                    k as usize
                } else {
                    cycle_start1 + ((k - cycle_start1 as i64) % cycle_len1 as i64) as usize
                };
                let g1 = path1[idx1];

                // (a-bi)^(e-k) を計算
                let exp2 = e - k;
                let idx2 = if exp2 < path2.len() as i64 {
                    exp2 as usize
                } else {
                    cycle_start2 + ((exp2 - cycle_start2 as i64) % cycle_len2 as i64) as usize
                };
                let g2 = path2[idx2];

                let g = g1.mul(g2, c);
                let state = (g.r * c + g.i) as usize;
                cnt[state] = (cnt[state] + count) % MOD;
            };

            let s = cycle_start1 as i64;
            let t = e - cycle_start2 as i64;

            if s > t {
                // 全探索
                for k in 0..=e {
                    add(k, 1);
                }
            } else {
                // 前端: [0, s-1]
                for k in 0..s {
                    add(k, 1);
                }
                // 後端: [t+1, e]
                for k in (t + 1)..=e {
                    add(k, 1);
                }
                // 中間: [s, t] → 周期性を利用して高速化
                let num = t - s + 1;
                if num > 0 {
                    let g = {
                        let mut x = cycle_len1 as i64;
                        let mut y = cycle_len2 as i64;
                        while y != 0 {
                            let r = x % y;
                            x = y;
                            y = r;
                        }
                        x
                    };
                    let lcm = (cycle_len1 as i64 / g) * cycle_len2 as i64;
                    for i in 0..lcm {
                        let times = num / lcm + if i < num % lcm { 1 } else { 0 };
                        if times > 0 {
                            add(s + i, times % MOD);
                        }
                    }
                }
            }

            for (state, &count) in cnt.iter().enumerate() {
                if count > 0 {
                    ways.push((state as i64, count));
                }
            }
        }

        // DP遷移: 現在の状態と新しい素因数の表現を畳み込む
        let mut nxt = vec![0i64; (c * c) as usize];
        for (state, &cur_cnt) in dp.iter().enumerate() {
            if cur_cnt == 0 {
                continue;
            }
            let cur_g = Gauss::new((state / c as usize) as i64, (state % c as usize) as i64, c);
            for &(w_state, w_cnt) in &ways {
                let w_g = Gauss::new((w_state / c) as i64, (w_state % c) as i64, c);
                let new_g = cur_g.mul(w_g, c);
                let new_state = (new_g.r * c + new_g.i) as usize;
                nxt[new_state] = (nxt[new_state] + cur_cnt * w_cnt) % MOD;
            }
        }
        dp = nxt;
    }

    // 条件: u ≡ -A (mod C), v ≡ -B (mod C)
    let target_r = ((-a % c) + c) % c;
    let target_i = ((-b % c) + c) % c;

    // 単位群の作用: (u, v) → (±u, ±v), (±v, ±u) の8通り
    let units = [
        Gauss::new(1, 0, c),
        Gauss::new(0, 1, c),
        Gauss::new(-1, 0, c),
        Gauss::new(0, -1, c),
    ];

    let mut ans = 0i64;
    for (state, &cnt) in dp.iter().enumerate() {
        if cnt == 0 {
            continue;
        }
        let g = Gauss::new((state / c as usize) as i64, (state % c as usize) as i64, c);
        for &unit in &units {
            let rotated = g.mul(unit, c);
            if rotated.r == target_r && rotated.i == target_i {
                ans = (ans + cnt) % MOD;
            }
        }
    }

    println!("{ans}");
}

Submission Info

Submission Time
Task G - Kyoen
User memoka
Language Rust (rustc 1.89.0)
Score 675
Code Size 7225 Byte
Status AC
Exec Time 1 ms
Memory 2244 KiB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 675 / 675
Status
AC × 4
AC × 58
Set Name Test Cases
Sample sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt
All hand_01.txt, hand_02.txt, random_01.txt, random_02.txt, random_03.txt, random_04.txt, random_05.txt, random_06.txt, random_07.txt, random_08.txt, random_09.txt, random_10.txt, random_11.txt, random_12.txt, random_13.txt, random_14.txt, random_15.txt, random_16.txt, random_17.txt, random_18.txt, random_19.txt, random_20.txt, random_21.txt, random_22.txt, random_23.txt, random_24.txt, random_25.txt, random_26.txt, random_27.txt, random_28.txt, random_29.txt, random_30.txt, random_31.txt, random_32.txt, random_33.txt, random_34.txt, random_35.txt, random_36.txt, random_37.txt, random_38.txt, random_39.txt, random_40.txt, random_41.txt, random_42.txt, random_43.txt, random_44.txt, random_45.txt, random_46.txt, random_47.txt, random_48.txt, random_49.txt, random_50.txt, random_51.txt, random_52.txt, sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt
Case Name Status Exec Time Memory
hand_01.txt AC 1 ms 2004 KiB
hand_02.txt AC 0 ms 1944 KiB
random_01.txt AC 0 ms 2044 KiB
random_02.txt AC 0 ms 2144 KiB
random_03.txt AC 0 ms 2004 KiB
random_04.txt AC 0 ms 1936 KiB
random_05.txt AC 0 ms 2096 KiB
random_06.txt AC 1 ms 2124 KiB
random_07.txt AC 0 ms 1940 KiB
random_08.txt AC 0 ms 2016 KiB
random_09.txt AC 0 ms 1840 KiB
random_10.txt AC 0 ms 2100 KiB
random_11.txt AC 0 ms 2108 KiB
random_12.txt AC 0 ms 1920 KiB
random_13.txt AC 0 ms 2048 KiB
random_14.txt AC 0 ms 1972 KiB
random_15.txt AC 1 ms 2144 KiB
random_16.txt AC 0 ms 2080 KiB
random_17.txt AC 0 ms 2048 KiB
random_18.txt AC 0 ms 1980 KiB
random_19.txt AC 1 ms 2108 KiB
random_20.txt AC 0 ms 2064 KiB
random_21.txt AC 1 ms 2048 KiB
random_22.txt AC 0 ms 2244 KiB
random_23.txt AC 1 ms 2028 KiB
random_24.txt AC 1 ms 2004 KiB
random_25.txt AC 0 ms 1960 KiB
random_26.txt AC 0 ms 2104 KiB
random_27.txt AC 0 ms 2068 KiB
random_28.txt AC 1 ms 2068 KiB
random_29.txt AC 1 ms 1964 KiB
random_30.txt AC 0 ms 2024 KiB
random_31.txt AC 1 ms 2108 KiB
random_32.txt AC 0 ms 2108 KiB
random_33.txt AC 1 ms 2072 KiB
random_34.txt AC 0 ms 2048 KiB
random_35.txt AC 0 ms 2016 KiB
random_36.txt AC 0 ms 1852 KiB
random_37.txt AC 0 ms 2048 KiB
random_38.txt AC 0 ms 2144 KiB
random_39.txt AC 0 ms 1988 KiB
random_40.txt AC 1 ms 1956 KiB
random_41.txt AC 0 ms 2160 KiB
random_42.txt AC 0 ms 2044 KiB
random_43.txt AC 1 ms 2068 KiB
random_44.txt AC 0 ms 2048 KiB
random_45.txt AC 0 ms 2048 KiB
random_46.txt AC 0 ms 1960 KiB
random_47.txt AC 0 ms 2068 KiB
random_48.txt AC 0 ms 1960 KiB
random_49.txt AC 0 ms 2004 KiB
random_50.txt AC 0 ms 2160 KiB
random_51.txt AC 0 ms 2048 KiB
random_52.txt AC 0 ms 2148 KiB
sample_01.txt AC 0 ms 1936 KiB
sample_02.txt AC 0 ms 1960 KiB
sample_03.txt AC 0 ms 2024 KiB
sample_04.txt AC 0 ms 2160 KiB