結果

問題 No.3416 マッチ棒パズル Extra
コンテスト
ユーザー akakimidori
提出日時 2025-12-25 00:29:47
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 833 ms / 3,000 ms
コード長 4,088 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 11,246 ms
コンパイル使用メモリ 400,248 KB
実行使用メモリ 7,848 KB
最終ジャッジ日時 2025-12-25 00:36:43
合計ジャッジ時間 24,123 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 26
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: function `naive` is never used
  --> src/main.rs:11:4
   |
11 | fn naive(n: usize) -> usize {
   |    ^^^^^
   |
   = note: `#[warn(dead_code)]` on by default

ソースコード

diff #
raw source code 

fn main() {
    input! {
        t: usize,
        ask: [usize; t],
    }
    for n in ask {
        println!("{}", fast(n) % 998244353);
    }
}

fn naive(n: usize) -> usize {
    let sq = (1usize..).find(|k| 2 * *k * (*k + 1) > n).unwrap() - 1;
    let v = (1..=sq)
        .map(|a| (n - a) / (2 * a + 1))
        .sum::<usize>();
    2 * v - sq * sq
}

fn fast(n: usize) -> u128 {
    let inside = |x, y| 2 * x * y + x + y >= n;
    let border = |x, y| 2 * x * y + x + y == n;
    let f = |x| 2 * x * (x + 1) <= n;
    let mut sq = (((2 * n + 1) as f64).sqrt().floor() as usize - 1) / 2;
    while !f(sq) {
        sq -= 1;
    }
    while f(sq + 1) {
        sq += 1;
    }
    let f = |x: usize| (n - x) / (2 * x + 1);
    let ceil_f = |x: usize| (n + x) / (2 * x + 1);

    // 打ち切り判定
    // 2xy + x + y = N
    // 2y + 2xy' + 1 + y' = 0
    // y' = -(2y + 1) / (2x + 1)
    // これを緩めた式でやってる
    let prun = |x, y, p: (usize, usize)| {
        (2 * y + 1) * p.0 >= (2 * x + 1) * p.1
    };
    let mut pos = (sq, ceil_f(sq));
    let mut stt = vec![((1, 0), (0, 1))];
    let mut ans = 0u128;
    while pos.0 > 3000000 {
        while stt.last().map_or(false, |&(_, p)| {
            let (x, y) = (pos.0 - p.0, pos.1 + p.1);
            !inside(x, y)
        }) {
            stt.pop();
        }
        let (mut l, mut r) = *stt.last().unwrap();
        loop {
            let d = (l.0 + r.0, l.1 + r.1);
            if d.0 >= pos.0 {
                break;
            }
            let (x, y) = (pos.0 - d.0, pos.1 + d.1);
            if inside(x, y) {
                stt.push((l, r));
                r = d;
            } else if prun(x, y, r) {
                break;
            } else {
                l = d;
            }
        }
        let (a, b) = r;
        let c = (a - 1) * (b - 1) / 2;
        if border(pos.0, pos.1) {
            ans += 1;
        }
        while inside(pos.0 - a, pos.1 + b) {
            ans += (pos.1 * a + c - 1) as u128;
            pos.1 += b;
            pos.0 -= a;
        }
    }
    ans += (1..=pos.0).map(|a| f(a as usize) as u128).sum::<u128>();
    2 * ans - (sq as u128).pow(2)
}

// ---------- begin floor sum ----------
// sum_{i = 0}^{n - 1} floor((ai + b) / m)
pub fn floor_sum(n: i128, m: i128, mut a: i128, mut b: i128) -> i128 {
    let mut ans = 0;
    const MOD: i128 = 998244353;
    ans += (a / m) % MOD * (n * (n - 1) / 2 % MOD) + (b / m) % MOD * n;
    a %= m;
    b %= m;
    let p = a * n + b;
    if p >= m {
        ans += floor_sum(p / m, a, m, p % m);
    }
    ans % MOD
}
// ---------- end floor sum ----------
// ---------- 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 ----------
0