結果

問題 No.2529 Treasure Hunter
ユーザー akakimidori
提出日時 2023-11-03 22:40:20
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 4 ms / 2,000 ms
コード長 15,044 bytes
コンパイル時間 14,919 ms
コンパイル使用メモリ 385,088 KB
実行使用メモリ 6,944 KB
最終ジャッジ日時 2024-09-25 21:04:31
合計ジャッジ時間 14,237 ms
ジャッジサーバーID
(参考情報)
judge2 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 1
other AC * 22
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: associated constant `IS_PRIME` is never used
   --> src/main.rs:324:11
    |
306 | impl<const M: u32> ConstantModulo<{ M }> {
    | ---------------------------------------- associated constant in this implementation
...
324 |     const IS_PRIME: () = assert!(is_prime(M));
    |           ^^^^^^^^
    |
    = note: `#[warn(dead_code)]` on by default

ソースコード

diff #
プレゼンテーションモードにする

use std::io::Write;
fn run() {
input! {
t: usize,
ask: [(usize, usize); t],
}
let out = std::io::stdout();
let mut out = std::io::BufWriter::new(out.lock());
for (n, m) in ask {
type Mat = Matrix<M, 3>;
let mut mat = Mat::zero();
for i in 0..3 {
for j in 0..3 {
let mat = &mut mat[i][j];
if j == 0 {
*mat = M::one();
continue;
}
if n < 4 && (i == 2 || j == 2) {
continue;
}
if j == 1 {
*mat = M::from(n);
} else {
*mat = M::from(n - 3 + (n - 3) * (n - 2) / 2);
}
if i == 1 {
if j == 1 {
*mat -= M::one();
} else {
*mat -= M::from(n - 3);
}
} else if i == 2 {
if j == 1 {
*mat -= M::new(2);
} else {
*mat -= M::from(2 * (n - 3)) - M::one();
}
}
}
}
let mut t = Mat::one();
let mut r = mat;
let mut m = m;
while m > 0 {
if m & 1 == 1 {
t = t * r;
}
r = r * r;
m >>= 1;
}
let ans = t[0].iter().fold(M::zero(), |s, a| s + *a);
writeln!(out, "{}", ans).ok();
}
}
fn main() {
run();
}
#[derive(Clone, Copy)]
pub struct Matrix<T, const K: usize>([[T; K]; K]);
impl<T, const K: usize> Matrix<T, K> {
fn new(a: [[T; K]; K]) -> Self {
Self(a)
}
}
impl<T, const K: usize> Zero for Matrix<T, K>
where
T: Zero + Copy,
{
fn zero() -> Self {
Self::new([[T::zero(); K]; K])
}
fn is_zero(&self) -> bool {
self.0.iter().flatten().all(|a| a.is_zero())
}
}
impl<T, const K: usize> Add for Matrix<T, K>
where
T: Zero + Copy,
{
type Output = Self;
fn add(self, rhs: Self) -> Self {
let mut res = Self::zero();
for ((res, a), b) in res.0.iter_mut().zip(self.0.iter()).zip(rhs.0.iter()) {
for ((res, a), b) in res.iter_mut().zip(a.iter()).zip(b.iter()) {
*res = *a + *b;
}
}
res
}
}
impl<T, const K: usize> One for Matrix<T, K>
where
T: Zero + One + Copy,
{
fn one() -> Self {
let mut res = Self::zero();
for (i, res) in res.0.iter_mut().enumerate() {
res[i] = T::one();
}
res
}
fn is_one(&self) -> bool {
self.0
.iter()
.enumerate()
.all(|(i, a)| a.iter().enumerate().all(|(j, a)| a.is_one() == (i == j)))
}
}
impl<T, const K: usize> Mul for Matrix<T, K>
where
T: Zero + One + Copy,
{
type Output = Self;
fn mul(self, rhs: Self) -> Self {
let mut res = Self::zero();
for (res, a) in res.0.iter_mut().zip(self.0.iter()) {
for (a, b) in a.iter().zip(rhs.0.iter()) {
for (res, b) in res.iter_mut().zip(b.iter()) {
*res = *res + *a * *b;
}
}
}
res
}
}
impl<T, const K: usize> Index<usize> for Matrix<T, K> {
type Output = [T; K];
fn index(&self, index: usize) -> &Self::Output {
&self.0[index]
}
}
impl<T, const K: usize> IndexMut<usize> for Matrix<T, K> {
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
&mut self.0[index]
}
}
// ---------- 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 ----------
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> {}
impl<T: Modulo> Zero for ModInt<T> {
fn zero() -> Self {
Self::zero()
}
fn is_zero(&self) -> bool {
self.is_zero()
}
}
impl<T: Modulo> One for ModInt<T> {
fn one() -> Self {
Self::one()
}
fn is_one(&self) -> bool {
self.is_one()
}
}
pub const fn pow_mod(mut r: u32, mut n: u32, m: u32) -> u32 {
let mut t = 1;
while n > 0 {
if n & 1 == 1 {
t = (t as u64 * r as u64 % m as u64) as u32;
}
r = (r as u64 * r as u64 % m as u64) as u32;
n >>= 1;
}
t
}
pub const fn primitive_root(p: u32) -> u32 {
let mut m = p - 1;
let mut f = [1; 30];
let mut k = 0;
let mut d = 2;
while d * d <= m {
if m % d == 0 {
f[k] = d;
k += 1;
}
while m % d == 0 {
m /= d;
}
d += 1;
}
if m > 1 {
f[k] = m;
k += 1;
}
let mut g = 1;
while g < p {
let mut ok = true;
let mut i = 0;
while i < k {
ok &= pow_mod(g, (p - 1) / f[i], p) > 1;
i += 1;
}
if ok {
break;
}
g += 1;
}
g
}
pub const fn is_prime(n: u32) -> bool {
if n <= 1 {
return false;
}
let mut d = 2;
while d * d <= n {
if n % d == 0 {
return false;
}
d += 1;
}
true
}
use std::marker::*;
use std::ops::*;
pub trait Modulo {
fn modulo() -> u32;
fn build(v: u32) -> u32;
fn reduce(v: u64) -> u32;
}
pub struct ConstantModulo<const M: u32>;
impl<const M: u32> ConstantModulo<{ M }> {
const ORDER: usize = (M - 1).trailing_zeros() as usize;
const PRIMITIVE_ROOT: u32 = primitive_root(M);
const ZETA: u32 = pow_mod(Self::PRIMITIVE_ROOT, (M - 1) >> Self::ORDER, M);
const REM: u32 = {
let mut t = 1u32;
let mut s = !M + 1;
let mut n = !0u32 >> 2;
while n > 0 {
if n & 1 == 1 {
t = t.wrapping_mul(s);
}
s = s.wrapping_mul(s);
n >>= 1;
}
t
};
const INI: u64 = ((1u128 << 64) % M as u128) as u64;
const IS_PRIME: () = assert!(is_prime(M));
}
impl<const M: u32> Modulo for ConstantModulo<{ M }> {
fn modulo() -> u32 {
M
}
fn build(v: u32) -> u32 {
Self::reduce(v as u64 * Self::INI)
}
fn reduce(x: u64) -> u32 {
debug_assert!(x < (Self::modulo() as u64) << 32);
let b = (x as u32 * Self::REM) as u64;
let t = x + b * M as u64;
let mut c = (t >> 32) as u32;
if c >= M {
c -= M;
}
c as u32
}
}
pub trait NTTFriendly {
fn order() -> usize;
fn zeta() -> u32;
}
impl<const M: u32> NTTFriendly for ConstantModulo<{ M }> {
fn order() -> usize {
Self::ORDER
}
fn zeta() -> u32 {
Self::ZETA
}
}
pub struct ModInt<T>(u32, PhantomData<fn() -> T>);
impl<T> Clone for ModInt<T> {
fn clone(&self) -> Self {
Self::build(self.0)
}
}
impl<T> Copy for ModInt<T> {}
impl<T: Modulo> Add for ModInt<T> {
type Output = Self;
fn add(self, rhs: Self) -> Self::Output {
let mut v = self.0 + rhs.0;
if v >= T::modulo() {
v -= T::modulo();
}
Self::build(v)
}
}
impl<T: Modulo> Sub for ModInt<T> {
type Output = Self;
fn sub(self, rhs: Self) -> Self::Output {
let mut v = self.0 - rhs.0;
if self.0 < rhs.0 {
v += T::modulo();
}
Self::build(v)
}
}
impl<T: Modulo> Mul for ModInt<T> {
type Output = Self;
fn mul(self, rhs: Self) -> Self::Output {
Self::build(T::reduce(self.0 as u64 * rhs.0 as u64))
}
}
impl<T: Modulo> AddAssign for ModInt<T> {
fn add_assign(&mut self, rhs: Self) {
*self = *self + rhs;
}
}
impl<T: Modulo> SubAssign for ModInt<T> {
fn sub_assign(&mut self, rhs: Self) {
*self = *self - rhs;
}
}
impl<T: Modulo> MulAssign for ModInt<T> {
fn mul_assign(&mut self, rhs: Self) {
*self = *self * rhs;
}
}
impl<T: Modulo> Neg for ModInt<T> {
type Output = Self;
fn neg(self) -> Self::Output {
if self.is_zero() {
Self::zero()
} else {
Self::build(T::modulo() - self.0)
}
}
}
impl<T: Modulo> std::fmt::Display for ModInt<T> {
fn fmt<'a>(&self, f: &mut std::fmt::Formatter<'a>) -> std::fmt::Result {
write!(f, "{}", self.get())
}
}
impl<T: Modulo> std::fmt::Debug for ModInt<T> {
fn fmt<'a>(&self, f: &mut std::fmt::Formatter<'a>) -> std::fmt::Result {
write!(f, "{}", self.get())
}
}
impl<T: Modulo> std::str::FromStr for ModInt<T> {
type Err = std::num::ParseIntError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let val = s.parse::<u32>()?;
Ok(ModInt::new(val))
}
}
impl<T: Modulo> From<usize> for ModInt<T> {
fn from(v: usize) -> Self {
Self::new_unchecked((v % T::modulo() as usize) as u32)
}
}
impl<T> ModInt<T> {
fn build(v: u32) -> Self {
ModInt(v, PhantomData)
}
pub fn is_zero(&self) -> bool {
self.0 == 0
}
}
impl<T: Modulo> ModInt<T> {
pub fn new_unchecked(v: u32) -> Self {
Self::build(T::build(v))
}
pub fn new(v: u32) -> Self {
Self::new_unchecked(v % T::modulo())
}
pub fn zero() -> Self {
Self::new_unchecked(0)
}
pub fn one() -> Self {
Self::new_unchecked(1)
}
pub fn get(&self) -> u32 {
T::reduce(self.0 as u64)
}
pub fn is_one(&self) -> bool {
self.get() == 1
}
pub fn pow(&self, mut n: u64) -> Self {
let mut t = Self::one();
let mut r = *self;
while n > 0 {
if n & 1 == 1 {
t *= r;
}
r *= r;
n >>= 1;
}
t
}
pub fn inv(&self) -> Self {
assert!(!self.is_zero());
self.pow((T::modulo() - 2) as u64)
}
pub fn fact(n: usize) -> Self {
(1..=n).fold(Self::one(), |s, a| s * Self::from(a))
}
}
pub trait ArrayAdd {
type Item;
fn add(&self, rhs: &[Self::Item]) -> Vec<Self::Item>;
}
impl<T> ArrayAdd for [T]
where
T: Zero + Copy,
{
type Item = T;
fn add(&self, rhs: &[Self::Item]) -> Vec<Self::Item> {
let mut c = vec![T::zero(); self.len().max(rhs.len())];
c[..self.len()].copy_from_slice(self);
c.add_assign(rhs);
c
}
}
pub trait ArrayAddAssign {
type Item;
fn add_assign(&mut self, rhs: &[Self::Item]);
}
impl<T> ArrayAddAssign for [T]
where
T: Add<Output = T> + Copy,
{
type Item = T;
fn add_assign(&mut self, rhs: &[Self::Item]) {
assert!(self.len() >= rhs.len());
self.iter_mut().zip(rhs).for_each(|(x, a)| *x = *x + *a);
}
}
impl<T> ArrayAddAssign for Vec<T>
where
T: Zero + Add<Output = T> + Copy,
{
type Item = T;
fn add_assign(&mut self, rhs: &[Self::Item]) {
if self.len() < rhs.len() {
self.resize(rhs.len(), T::zero());
}
self.as_mut_slice().add_assign(rhs);
}
}
pub trait ArraySub {
type Item;
fn sub(&self, rhs: &[Self::Item]) -> Vec<Self::Item>;
}
impl<T> ArraySub for [T]
where
T: Zero + Sub<Output = T> + Copy,
{
type Item = T;
fn sub(&self, rhs: &[Self::Item]) -> Vec<Self::Item> {
let mut c = vec![T::zero(); self.len().max(rhs.len())];
c[..self.len()].copy_from_slice(self);
c.sub_assign(rhs);
c
}
}
pub trait ArraySubAssign {
type Item;
fn sub_assign(&mut self, rhs: &[Self::Item]);
}
impl<T> ArraySubAssign for [T]
where
T: Sub<Output = T> + Copy,
{
type Item = T;
fn sub_assign(&mut self, rhs: &[Self::Item]) {
assert!(self.len() >= rhs.len());
self.iter_mut().zip(rhs).for_each(|(x, a)| *x = *x - *a);
}
}
impl<T> ArraySubAssign for Vec<T>
where
T: Zero + Sub<Output = T> + Copy,
{
type Item = T;
fn sub_assign(&mut self, rhs: &[Self::Item]) {
if self.len() < rhs.len() {
self.resize(rhs.len(), T::zero());
}
self.as_mut_slice().sub_assign(rhs);
}
}
pub trait ArrayDot {
type Item;
fn dot(&self, rhs: &[Self::Item]) -> Vec<Self::Item>;
}
impl<T> ArrayDot for [T]
where
T: Mul<Output = T> + Copy,
{
type Item = T;
fn dot(&self, rhs: &[Self::Item]) -> Vec<Self::Item> {
assert!(self.len() == rhs.len());
self.iter().zip(rhs).map(|p| *p.0 * *p.1).collect()
}
}
pub trait ArrayDotAssign {
type Item;
fn dot_assign(&mut self, rhs: &[Self::Item]);
}
impl<T> ArrayDotAssign for [T]
where
T: MulAssign + Copy,
{
type Item = T;
fn dot_assign(&mut self, rhs: &[Self::Item]) {
assert!(self.len() == rhs.len());
self.iter_mut().zip(rhs).for_each(|(x, a)| *x *= *a);
}
}
pub trait ArrayMul {
type Item;
fn mul(&self, rhs: &[Self::Item]) -> Vec<Self::Item>;
}
impl<T> ArrayMul for [T]
where
T: Zero + Mul<Output = T> + Copy,
{
type Item = T;
fn mul(&self, rhs: &[Self::Item]) -> Vec<Self::Item> {
if self.is_empty() || rhs.is_empty() {
return vec![];
}
let mut res = vec![T::zero(); self.len() + rhs.len() - 1];
for (i, a) in self.iter().enumerate() {
for (c, b) in res[i..].iter_mut().zip(rhs) {
*c = *c + *a * *b;
}
}
res
}
}
const PRIME: u32 = 998_244_353;
type S = ConstantModulo<PRIME>;
type M = ModInt<S>;
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0