#938811 (Rust) No.2589 Prepare Integers

• 問題: No.2589 Prepare Integers
• ユーザー: akakimidori
• 提出日時: 2023-12-17 08:40:04
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 70 ms / 2,000 ms
• コード長: 12,296 bytes
• コンパイル時間: 14,613 ms
• コンパイル使用メモリ: 378,732 KB
• 実行使用メモリ: 5,376 KB
• 最終ジャッジ日時: 2024-09-27 07:58:40
• 合計ジャッジ時間: 15,484 ms
• ジャッジサーバーID （参考情報）: judge4 / judge2

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 5,248 KB
testcase_01	AC	1 ms; 5,376 KB
testcase_02	AC	1 ms; 5,376 KB
testcase_03	AC	1 ms; 5,376 KB
testcase_04	AC	70 ms; 5,376 KB
testcase_05	AC	39 ms; 5,376 KB
testcase_06	AC	28 ms; 5,376 KB
testcase_07	AC	26 ms; 5,376 KB
testcase_08	AC	19 ms; 5,376 KB
testcase_09	AC	16 ms; 5,376 KB
testcase_10	AC	16 ms; 5,376 KB
testcase_11	AC	17 ms; 5,376 KB
testcase_12	AC	32 ms; 5,376 KB
testcase_13	AC	27 ms; 5,376 KB
testcase_14	AC	24 ms; 5,376 KB
testcase_15	AC	24 ms; 5,376 KB
testcase_16	AC	23 ms; 5,376 KB
testcase_17	AC	18 ms; 5,376 KB
testcase_18	AC	16 ms; 5,376 KB
testcase_19	AC	16 ms; 5,376 KB
testcase_20	AC	17 ms; 5,376 KB
testcase_21	AC	16 ms; 5,376 KB
testcase_22	AC	16 ms; 5,376 KB
testcase_23	AC	17 ms; 5,376 KB
testcase_24	AC	17 ms; 5,376 KB
testcase_25	AC	17 ms; 5,376 KB
testcase_26	AC	17 ms; 5,376 KB
testcase_27	AC	15 ms; 5,376 KB

コンパイルメッセージ

warning: function `stress` is never used
  --> src/main.rs:28:4
   |
28 | fn stress() {
   |    ^^^^^^
   |
   = note: `#[warn(dead_code)]` on by default

warning: function `rand_memory` is never used
   --> src/main.rs:451:4
    |
451 | fn rand_memory() -> usize {
    |    ^^^^^^^^^^^

warning: function `rand` is never used
   --> src/main.rs:455:4
    |
455 | fn rand() -> usize {
    |    ^^^^

warning: function `shuffle` is never used
   --> src/main.rs:468:4
    |
468 | fn shuffle<T>(a: &mut [T]) {
    |    ^^^^^^^

ソースコード

// 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() {
    //stress();
    run();
}

fn stress() {
    for t in 0..100000 {
        if t % 100 == 0 {
            println!("{}", t);
        }
        let m = (rand() % 10 + 2) as i64;
        let mut space = std::collections::BTreeSet::new();
        space.insert(0);
        let add = |mut a: i64, mut b: i64| -> i64 {
            let mut d = vec![];
            while a > 0 || b > 0 {
                d.push((a + b) % m);
                a /= m;
                b /= m;
            }
            d.iter().rfold(0, |s, a| s * m + *a)
        };
        let mut base = BasisModM::new(m);
        let mut memo = vec![];
        for _ in 0..100 {
            let op = rand() % 3 + 1;
            let p = (rand() % 1000 + 1) as i64;
            if op == 1 {
                base.add(p);
                memo.push(p);
                if space.insert(p) {
                    while {
                        let src = space.clone();
                        let mut cond = false;
                        for u in src {
                            cond |= space.insert(add(u, p));
                        }
                        cond
                    } {}
                }
            } else if op == 2 {
                let x = base.kth_term(p - 1);
                assert_eq!(
                    x,
                    space.iter().nth(p as usize - 1).cloned(),
                    "ERROR: {}, {:?}",
                    m,
                    memo
                );
            } else {
                let cnt = base.count_less(p);
                let ans = space.iter().take_while(|u| **u <= p).count() as i64;
                assert_eq!(cnt, ans, "ERROR: {}, {:?}", m, memo);
            }
        }
    }
}

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;
        }
        self.add_inner(b);
    }
    // 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 up = vec![0; len];
        for d in up.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(up.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];
            } else if now[i] > *d {
                return ans;
            }
        }
        ans + 1
    }
    fn add_inner(&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 {
                let 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 g = self.base[i][i];
                let x = w[i] / g;
                let y = b[i] / g;
                for (v, w) in b.iter_mut().zip(w.iter()) {
                    *v = (x * *v - y * *w + m * m) % m;
                }
            }
        }
    }
}

// ---------- 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
    }
}

fn rand_memory() -> usize {
    Box::into_raw(Box::new("I hope this is a random number")) as usize
}

fn rand() -> usize {
    static mut X: usize = 0;
    unsafe {
        if X == 0 {
            X = rand_memory();
        }
        X ^= X << 13;
        X ^= X >> 17;
        X ^= X << 5;
        X
    }
}

fn shuffle<T>(a: &mut [T]) {
    for i in 1..a.len() {
        let p = rand() % (i + 1);
        a.swap(i, p);
    }
}