#938806 (Rust) No.2589 Prepare Integers

• 問題: No.2589 Prepare Integers
• ユーザー: akakimidori
• 提出日時: 2023-12-17 07:24:45
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 11,474 bytes
• コンパイル時間: 13,018 ms
• コンパイル使用メモリ: 377,112 KB
• 実行使用メモリ: 6,944 KB
• 最終ジャッジ日時: 2024-09-27 07:57:39
• 合計ジャッジ時間: 14,842 ms
• ジャッジサーバーID （参考情報）: judge3 / judge1

ファイルパターン	結果
sample	AC * 4
other	AC * 9 WA * 15

ソースコード

// https://codeforces.com/blog/entry/98376
// これを読解する

fn run() {
    input! {
        m: i64,
        q: usize,
        ask: [(u8, i64); q],
    }
    let mut base = BasisModM::new(m);
    for (t, x) in ask {
        if t == 1 {
            base.add(x);
        } else if t == 2 {
            let ans = base.kth_term(x - 1).unwrap_or(-1);
            println!("{}", ans);
        } else {
            println!("{}", base.count_less(x));
        }
    }
}

fn main() {
    run();
}

fn gcd(a: i64, b: i64) -> i64 {
    if b == 0 {
        a
    } else {
        gcd(b, a % b)
    }
}

fn inv(a: i64, b: i64) -> i64 {
    if a == 1 {
        1
    } else {
        let c = inv(b % a, a);
        b + (-b * c + 1) / a
    }
}

fn ext_gcd(a: i64, b: i64) -> (i64, i64) {
    let g = gcd(a, b);
    let (mut a, mut b) = (a / g, b / g);
    type Mat = Matrix<i64, 2, 2>;
    let mut mat = Mat::one();
    while a > 0 && b > 0 {
        let mut p = Mat::one();
        if a >= b {
            let d = a / b;
            a -= d * b;
            p[1][0] = -d;
        } else {
            let d = b / a;
            b -= d * a;
            p[0][1] = -d;
        }
        mat = mat * p;
    }
    if a > 0 {
        (mat[0][0], mat[1][0])
    } else {
        (mat[0][1], mat[1][1])
    }
}

pub struct BasisModM {
    m: i64,
    len: usize,
    base: Vec<Vec<i64>>,
}

const UP: i64 = 1000000000;

impl BasisModM {
    fn new(m: i64) -> Self {
        assert!(m > 1);
        let mut up = UP;
        let mut len = 0;
        while up > 0 {
            up /= m;
            len += 1;
        }
        Self {
            m,
            len,
            base: vec![vec![0; len]; len],
        }
    }
    fn add(&mut self, mut v: i64) {
        let mut b = vec![0; self.len];
        for b in b.iter_mut() {
            *b = v % self.m;
            v /= self.m;
        }
        //println!("add {:?}", b);
        self.add_rec(b);
        let m = self.m;
        for i in (0..self.len).rev() {
            if self.base[i][i] == 0 {
                continue;
            }
            let src = std::mem::take(&mut self.base[i]);
            for base in self.base.iter_mut().skip(i + 1) {
                if base[i] >= src[i] {
                    let d = base[i] / src[i];
                    for (base, src) in base.iter_mut().zip(src.iter()) {
                        *base = (*base - d * *src + m * m) % m;
                    }
                }
            }
            self.base[i] = src;
        }
        /*
        for b in self.base.iter() {
            println!("{:?}", b);
        }
        println!();
        */
    }
    // 0-indexed
    fn kth_term(&self, x: i64) -> Option<i64> {
        let mut x = x;
        let mut d = vec![0; self.len];
        for (i, (d, base)) in d.iter_mut().zip(self.base.iter()).enumerate() {
            if base[i] != 0 {
                let p = self.m / base[i];
                *d = x % p;
                x /= p;
            }
        }
        if x > 0 {
            return None;
        }
        let m = self.m;
        let mut ans = vec![0; self.len];
        for (i, (d, base)) in d.iter_mut().zip(self.base.iter()).enumerate().rev() {
            if base[i] == 0 {
                continue;
            }
            let sub = ans[i] / base[i];
            for (ans, b) in ans.iter_mut().zip(base.iter()) {
                *ans = (*ans + (*d - sub) * *b + m * m) % m;
            }
        }
        Some(ans.iter().rfold(0, |s, a| s * m + *a))
    }
    // x 以下であるものの個数を求める
    fn count_less(&self, mut x: i64) -> i64 {
        let m = self.m;
        let base = &self.base;
        let len = self.len;
        let mut d = vec![0; len];
        for d in d.iter_mut() {
            *d = x % m;
            x /= m;
        }
        let mut dp = vec![1];
        let mut prod = 1;
        for (i, base) in base.iter().enumerate() {
            if base[i] != 0 {
                prod *= m / base[i];
            }
            dp.push(prod);
        }
        if x > 0 {
            return prod;
        }
        let mut ans = 0;
        let mut now = vec![0; len];
        for (i, (base, d)) in base.iter().zip(d.iter()).enumerate().rev() {
            if base[i] != 0 {
                let sub = now[i] / base[i];
                for (now, base) in now.iter_mut().zip(base.iter()) {
                    *now = (*now - sub * *base + m * m) % m;
                }
                if *d >= now[i] {
                    let q = (*d - now[i] + base[i] - 1) / base[i];
                    ans += q * dp[i];
                    for (now, base) in now.iter_mut().zip(base.iter()) {
                        *now = (*now + q * *base) % m;
                    }
                }
                if now[i] > *d {
                    return ans;
                }
            } else if now[i] < *d {
                return ans + dp[i + 1];
            }
        }
        ans + 1
    }
    fn add_rec(&mut self, mut b: Vec<i64>) {
        let m = self.m;
        for i in (0..self.len).rev() {
            if b[i] == 0 {
                continue;
            }
            if self.base[i][i] == 0 {
                let g = gcd(m, b[i]);
                let mul = inv(b[i] / g, m / g);
                for (base, b) in self.base[i].iter_mut().zip(b.iter_mut()) {
                    *base = *b * mul % m;
                    *b = *b * (m / g) % m;
                }
            } else if b[i] % self.base[i][i] == 0 {
                let mul = b[i] / self.base[i][i];
                for (base, b) in self.base[i].iter().zip(b.iter_mut()) {
                    *b = (*b - mul * *base + m * m) % m;
                }
            } else {
                let mut w = self.base[i].clone();
                let (s, t) = ext_gcd(b[i], w[i]);
                let (s, t) = ((s + m) % m, (t + m) % m);
                for ((b, w), v) in self.base[i].iter_mut().zip(b.iter()).zip(w.iter()) {
                    *b = (*w * s + *v * t) % m;
                }
                let mul = b[i] / self.base[i][i];
                for (v, a) in b.iter_mut().zip(self.base[i].iter()) {
                    *v = (*v - mul * *a + m * m) % m;
                }
                let mul = w[i] / self.base[i][i];
                for (w, b) in w.iter_mut().zip(self.base[i].iter()) {
                    *w = (*w - mul * *b + m * m) % m;
                }
                self.add_rec(w);
                let mut a = self.base[i].clone();
                let mul = m / a[i];
                for a in a.iter_mut() {
                    *a = *a * mul % m;
                }
                self.add_rec(a);
            }
        }
    }
}

// ---------- 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 ----------
use std::ops::*;

pub trait Zero: Sized + Add<Self, Output = Self> {
    fn zero() -> Self;
    fn is_zero(&self) -> bool;
}

pub trait One: Sized + Mul<Self, Output = Self> {
    fn one() -> Self;
    fn is_one(&self) -> bool;
}

pub trait Ring: Zero + One + Sub<Output = Self> {}

pub trait Field: Ring + Div<Output = Self> {}

#[derive(Clone, Copy)]
pub struct Matrix<T, const R: usize, const C: usize>([[T; C]; R]);

impl<T, const R: usize, const C: usize> Matrix<T, R, C> {
    pub fn new(a: [[T; C]; R]) -> Self {
        Self(a)
    }
}

impl<T, const R: usize, const C: usize> Add for Matrix<T, R, C>
where
    T: Add<Output = T> + Copy,
{
    type Output = Self;
    fn add(self, rhs: Self) -> Self::Output {
        let mut res = self.0;
        for (res, b) in res.iter_mut().flatten().zip(rhs.0.into_iter().flatten()) {
            *res = *res + b;
        }
        Matrix::new(res)
    }
}

impl<T, const R: usize, const C: usize> Sub for Matrix<T, R, C>
where
    T: Sub<Output = T> + Copy,
{
    type Output = Self;
    fn sub(self, rhs: Self) -> Self::Output {
        let mut res = self.0;
        for (res, b) in res.iter_mut().flatten().zip(rhs.0.into_iter().flatten()) {
            *res = *res - b;
        }
        Matrix::new(res)
    }
}

impl<T, const R: usize, const M: usize, const C: usize> Mul<Matrix<T, M, C>> for Matrix<T, R, M>
where
    T: Zero + Mul<Output = T> + Copy,
{
    type Output = Matrix<T, R, C>;
    fn mul(self, rhs: Matrix<T, M, C>) -> Self::Output {
        let mut c = Matrix::new([[T::zero(); C]; R]);
        for (c, a) in c.0.iter_mut().zip(self.0.iter()) {
            for (a, b) in a.iter().zip(rhs.0.iter()) {
                for (c, b) in c.iter_mut().zip(b.iter()) {
                    *c = *c + *a * *b;
                }
            }
        }
        c
    }
}

impl<T, const R: usize, const C: usize> Zero for Matrix<T, R, C>
where
    T: Zero + Copy,
{
    fn zero() -> Self {
        Self::new([[T::zero(); C]; R])
    }
    fn is_zero(&self) -> bool {
        self.0.iter().flatten().all(|p| p.is_zero())
    }
}

impl<T, const N: usize> One for Matrix<T, N, N>
where
    T: Zero + One + Copy,
{
    fn one() -> Self {
        let mut res = Self::new([[T::zero(); N]; N]);
        for (i, res) in res.0.iter_mut().enumerate() {
            res[i] = T::one();
        }
        res
    }
    fn is_one(&self) -> bool {
        for (i, res) in self.0.iter().enumerate() {
            for (j, res) in res.iter().enumerate() {
                if i == j && !res.is_one() {
                    return false;
                }
                if i != j && !res.is_zero() {
                    return false;
                }
            }
        }
        true
    }
}

impl<T, const R: usize, const C: usize> Index<usize> for Matrix<T, R, C> {
    type Output = [T; C];
    fn index(&self, x: usize) -> &Self::Output {
        &self.0[x]
    }
}

impl<T, const R: usize, const C: usize> IndexMut<usize> for Matrix<T, R, C> {
    fn index_mut(&mut self, x: usize) -> &mut Self::Output {
        &mut self.0[x]
    }
}

impl Zero for i64 {
    fn zero() -> Self {
        0
    }
    fn is_zero(&self) -> bool {
        *self == 0
    }
}

impl One for i64 {
    fn one() -> Self {
        1
    }
    fn is_one(&self) -> bool {
        *self == 1
    }
}