#1099174 (Rust) No.3182 recurrence relation’s intersection sum

• 問題: No.3182 recurrence relation’s intersection sum
• ユーザー: Moss_Local
• 提出日時: 2025-06-13 23:49:56
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 22 ms / 2,000 ms
• コード長: 9,125 bytes
• コンパイル時間: 12,219 ms
• コンパイル使用メモリ: 399,224 KB
• 実行使用メモリ: 7,716 KB
• 最終ジャッジ日時: 2025-10-15 12:37:36
• 合計ジャッジ時間: 14,956 ms
• ジャッジサーバーID （参考情報）: judge5 / judge1

ファイルパターン	結果
sample	AC * 3
other	AC * 40

コンパイルメッセージ

warning: variable does not need to be mutable
   --> src/main.rs:109:9
    |
109 |     let mut vec: Vec<i64> = read_vec();
    |         ----^^^
    |         |
    |         help: remove this `mut`
    |
    = note: `#[warn(unused_mut)]` on by default

warning: variable does not need to be mutable
   --> src/main.rs:115:9
    |
115 |     let mut vec: Vec<i64> = read_vec();
    |         ----^^^
    |         |
    |         help: remove this `mut`

warning: variable does not need to be mutable
   --> src/main.rs:120:9
    |
120 |     let mut vec: Vec<usize> = read_vec();
    |         ----^^^
    |         |
    |         help: remove this `mut`

warning: variable does not need to be mutable
   --> src/main.rs:125:9
    |
125 |     let mut vec: Vec<usize> = read_vec();
    |         ----^^^
    |         |
    |         help: remove this `mut`

warning: variable `N` should have a snake case name
   --> src/main.rs:364:9
    |
364 |     let N = 500;
    |         ^ help: convert the identifier to snake case: `n`
    |
    = note: `#[warn(non_snake_case)]` on by default

ソースコード

// -*- coding:utf-8-unix -*-
// #![feature(map_first_last)]
#![allow(dead_code)]
#![allow(unused_imports)]
#![allow(unused_macros)]

use std::cmp::*;
use std::collections::*;
use std::fmt::*;
use std::hash::*;
use std::io::BufRead;
use std::iter::FromIterator;
use std::*;

const INF: i64 = 1223372036854775807;
const UINF: usize = INF as usize;
const LINF: i64 = 2147483647;
const INF128: i128 = 1223372036854775807000000000000;
const MOD1: i64 = 1000000007;
const MOD9: i64 = 998244353;
const MOD: i64 = MOD9;
const UMOD: usize = MOD as usize;
const M_PI: f64 = 3.14159265358979323846;

macro_rules! p {
    ($x:expr) => {
        //if expr
        println!("{}", $x);
    };
}

macro_rules! vp {
    // vector print separate with space
    ($x:expr) => {
        println!(
            "{}",
            $x.iter()
                .map(|x| x.to_string())
                .collect::<Vec<_>>()
                .join(" ")
        );
    };
}

macro_rules! d {
    ($x:expr) => {
        eprintln!("{:?}", $x);
    };
}
macro_rules! yn {
    ($val:expr) => {
        if $val {
            println!("Yes");
        } else {
            println!("No");
        }
    };
}

macro_rules! map{
    // declear btreemap
    ($($key:expr => $val:expr),*) => {
        {
            let mut map = ::std::collections::BTreeMap::new();
            $(
                map.insert($key, $val);
            )*
            map
        }
    };
}

macro_rules! set{
    // declear btreemap
    ($($key:expr),*) => {
        {
            let mut set = ::std::collections::BTreeSet::new();
            $(
                set.insert($key);
            )*
            set
        }
    };
}

//input output
#[allow(dead_code)]
fn read<T: std::str::FromStr>() -> T {
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).ok();
    s.trim().parse().ok().unwrap()
}

#[allow(dead_code)]
fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>()
        .split_whitespace()
        .map(|e| e.parse().ok().unwrap())
        .collect()
}

#[allow(dead_code)]
fn read_mat<T: std::str::FromStr>(n: u32) -> Vec<Vec<T>> {
    (0..n).map(|_| read_vec()).collect()
}

#[allow(dead_code)]
fn readii() -> (i64, i64) {
    let mut vec: Vec<i64> = read_vec();
    (vec[0], vec[1])
}

#[allow(dead_code)]
fn readiii() -> (i64, i64, i64) {
    let mut vec: Vec<i64> = read_vec();
    (vec[0], vec[1], vec[2])
}
#[allow(dead_code)]
fn readuu() -> (usize, usize) {
    let mut vec: Vec<usize> = read_vec();
    (vec[0], vec[1])
}

fn readuuu() -> (usize, usize, usize) {
    let mut vec: Vec<usize> = read_vec();
    (vec[0], vec[1], vec[2])
}

const PRIMITIVE_ROOT: usize = 3;

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
struct ModInt(usize);

impl ModInt {
    fn new(n: usize) -> Self {
        ModInt(n % UMOD)
    }
    fn zero() -> Self {
        ModInt(0)
    }
    fn one() -> Self {
        ModInt(1)
    }
    fn pow(self, mut e: u64) -> Self {
        let mut base = self;
        let mut res = ModInt::one();
        while e > 0 {
            if e & 1 == 1 {
                res *= base;
            }
            base *= base;
            e >>= 1;
        }
        res
    }
    fn inv(self) -> Self {
        // Fermat's little theorem: a^(UMOD-2)
        self.pow((UMOD as u64) - 2)
    }
}

impl From<u64> for ModInt {
    fn from(x: u64) -> Self {
        ModInt((x % (UMOD as u64)) as usize)
    }
}

use std::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};

impl Add for ModInt {
    type Output = ModInt;
    fn add(self, rhs: ModInt) -> ModInt {
        let sum = self.0 + rhs.0;
        ModInt(if sum >= UMOD { sum - UMOD } else { sum })
    }
}

impl Sub for ModInt {
    type Output = ModInt;
    fn sub(self, rhs: ModInt) -> ModInt {
        ModInt(if self.0 < rhs.0 {
            self.0 + UMOD - rhs.0
        } else {
            self.0 - rhs.0
        })
    }
}

impl Mul for ModInt {
    type Output = ModInt;
    fn mul(self, rhs: ModInt) -> ModInt {
        ModInt((self.0 * rhs.0) % UMOD)
    }
}

impl AddAssign for ModInt {
    fn add_assign(&mut self, rhs: ModInt) {
        *self = *self + rhs;
    }
}
impl SubAssign for ModInt {
    fn sub_assign(&mut self, rhs: ModInt) {
        *self = *self - rhs;
    }
}
impl MulAssign for ModInt {
    fn mul_assign(&mut self, rhs: ModInt) {
        *self = *self * rhs;
    }
}

// Number Theoretic Transform (NTT)
fn bit_reverse(a: &mut [ModInt]) {
    let n = a.len();
    let mut j = 0;
    for i in 1..n {
        let mut bit = n >> 1;
        while j & bit != 0 {
            j ^= bit;
            bit >>= 1;
        }
        j |= bit;
        if i < j {
            a.swap(i, j);
        }
    }
}

fn ntt(a: &mut [ModInt], invert: bool) {
    let n = a.len();
    bit_reverse(a);
    let mut len = 2;
    while len <= n {
        let wlen = ModInt::from(PRIMITIVE_ROOT as u64).pow((UMOD as u64 - 1) / len as u64);
        let wlen = if invert { wlen.inv() } else { wlen };
        for i in (0..n).step_by(len) {
            let mut w = ModInt::one();
            for j in 0..len / 2 {
                let u = a[i + j];
                let v = a[i + j + len / 2] * w;
                a[i + j] = u + v;
                a[i + j + len / 2] = u - v;
                w *= wlen;
            }
        }
        len <<= 1;
    }
    if invert {
        let inv_n = ModInt::from(n as u64).inv();
        for x in a.iter_mut() {
            *x *= inv_n;
        }
    }
}

fn convolution(a: &[ModInt], b: &[ModInt]) -> Vec<ModInt> {
    let mut n = 1;
    while n < a.len() + b.len() - 1 {
        n <<= 1;
    }
    let mut fa = a.to_vec();
    fa.resize(n, ModInt::zero());
    let mut fb = b.to_vec();
    fb.resize(n, ModInt::zero());
    ntt(&mut fa, false);
    ntt(&mut fb, false);
    for i in 0..n {
        fa[i] *= fb[i];
    }
    ntt(&mut fa, true);
    fa.resize(a.len() + b.len() - 1, ModInt::zero());
    fa
}

// Berlekamp–Massey algorithm
fn berlekamp_massey(s: &[ModInt]) -> Vec<ModInt> {
    let n = s.len();
    let mut c = vec![ModInt::one()];
    let mut b = vec![ModInt::one()];
    let mut l = 0;
    let mut m = 1;
    let mut bb = ModInt::one();
    for i in 0..n {
        let mut d = ModInt::zero();
        for j in 0..=l {
            d += c[j] * s[i - j];
        }
        if d == ModInt::zero() {
            m += 1;
        } else if 2 * l <= i {
            let t = c.clone();
            let coef = d * bb.inv();
            c.resize((b.len() + m).max(c.len()), ModInt::zero());
            for j in 0..b.len() {
                c[j + m] -= coef * b[j];
            }
            l = i + 1 - l;
            b = t;
            bb = d;
            m = 1;
        } else {
            let coef = d * bb.inv();
            c.resize((b.len() + m).max(c.len()), ModInt::zero());
            for j in 0..b.len() {
                c[j + m] -= coef * b[j];
            }
            m += 1;
        }
    }
    c.remove(0);
    for x in c.iter_mut() {
        *x = ModInt::zero() - *x;
    }
    c
}

// Bostan–Mori to compute [z^N] num(z)/den(z)
fn bostan_mori(mut num: Vec<ModInt>, mut den: Vec<ModInt>, mut n: i64) -> ModInt {
    while n > 0 {
        let mut den_neg = den.clone();
        for (i, x) in den_neg.iter_mut().enumerate() {
            if i % 2 == 1 {
                *x = ModInt::zero() - *x;
            }
        }
        let f2 = convolution(&num, &den_neg);
        let g2 = convolution(&den, &den_neg);
        let num2: Vec<ModInt> = if n % 2 == 0 {
            f2.iter().step_by(2).cloned().collect()
        } else {
            f2.iter().skip(1).step_by(2).cloned().collect()
        };
        let den2: Vec<ModInt> = g2.iter().step_by(2).cloned().collect();
        num = num2;
        den = den2;
        n /= 2;
    }
    num[0] * den[0].inv()
}

// Compute k-th term (0-based) of linearly recurrent sequence with initial a and recurrence c
fn linear_recurrence(a: &[ModInt], c: &[ModInt], k: i64) -> ModInt {
    let n = c.len();
    if n == 0 {
        return ModInt::zero();
    }
    // D(z) = 1 - sum_{i=0..n) c[i] z^{i+1}
    let mut dnm = vec![ModInt::one(); 1 + n];
    for i in 0..n {
        dnm[i + 1] = ModInt::zero() - c[i];
    }
    let mut a_vec = a.to_vec();
    a_vec.resize(n, ModInt::zero());
    let num = convolution(&dnm, &a_vec)[..n].to_vec();
    bostan_mori(num, dnm, k)
}

fn main() {
    let (k, l, r) = readuuu();
    let k = k as i64;
    let l = l as i64;

    // Precompute sequence up to N
    let N = 500;
    let mut a = vec![ModInt::zero(); N + 1];
    a[0] = ModInt::one();
    for n in 0..N {
        let term = ModInt::from(k as u64) * a[n]
            + ModInt::from(n as u64).pow(k as u64)
            + ModInt::from(k as u64).pow(n as u64);
        a[n + 1] = term;
    }
    // Prefix sums b
    let mut b = vec![ModInt::zero(); N + 2];
    for i in 0..=N {
        b[i + 1] = b[i] + a[i];
    }
    // Find minimal linear recurrence for b
    let c = berlekamp_massey(&b);
    let res_r = linear_recurrence(&b, &c, (r + 1) as i64);
    let res_l = linear_recurrence(&b, &c, l as i64);
    let ans = res_r - res_l;
    println!("{}", ans.0);
}