結果
問題 |
No.2611 Count 01
|
ユーザー |
![]() |
提出日時 | 2025-08-17 12:03:37 |
言語 | Rust (1.83.0 + proconio) |
結果 |
AC
|
実行時間 | 718 ms / 6,000 ms |
コード長 | 27,240 bytes |
コンパイル時間 | 12,780 ms |
コンパイル使用メモリ | 393,952 KB |
実行使用メモリ | 55,412 KB |
最終ジャッジ日時 | 2025-08-17 12:04:07 |
合計ジャッジ時間 | 29,848 ms |
ジャッジサーバーID (参考情報) |
judge4 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 20 |
コンパイルメッセージ
warning: type alias `Map` is never used --> src/main.rs:80:6 | 80 | type Map<K, V> = BTreeMap<K, V>; | ^^^ | = note: `#[warn(dead_code)]` on by default warning: type alias `Set` is never used --> src/main.rs:81:6 | 81 | type Set<T> = BTreeSet<T>; | ^^^ warning: type alias `Deque` is never used --> src/main.rs:82:6 | 82 | type Deque<T> = VecDeque<T>; | ^^^^^
ソースコード
fn run<W: Write>(sc: &mut scanner::Scanner, out: &mut std::io::BufWriter<W>) { let n: usize = sc.next(); let q: usize = sc.next(); let mut s = sc.next_bytes(); type M = ModInt<998244353>; const N: usize = 4; type T = (Matrix<M, N, N>, Matrix<M, 1, N>, Matrix<M, N, 1>, Matrix<M, 1, 1>); let mut ini = vec![]; let mut trans = Matrix::<M, N, N>::one(); let mut l = Matrix::<M, 1, N>::zero(); let mut r = Matrix::<M, N, 1>::zero(); let sum = Matrix::<M, 1, 1>::zero(); trans[0][1] = M::one(); trans[2][3] = M::one(); l[0][0] = M::one(); r[3][0] = M::one(); ini.push((trans, l * trans, trans * r, sum)); trans = one(); trans[0][2] = M::one(); trans[1][3] = M::one(); ini.push((trans, l * trans, trans * r, sum)); let mut seg = SegmentTreePURQ::new(n, (one(), zero(), zero(), zero()), |a: &T, b| { let trans = a.0 * b.0; let l = a.1 * b.0 + b.1; let r = a.2 + a.0 * b.2; let s = a.3 + b.3 + a.1 * b.2; (trans, l, r, s) }); for (i, s) in s.iter().enumerate() { seg.update_tmp(i, ini[(*s - b'0') as usize]); } seg.update_all(); for _ in 0..q { let op: u8 = sc.next(); if op == 1 { let x = sc.next::<usize>() - 1; s[x] ^= b'0' ^ b'1'; seg.update(x, ini[(s[x] - b'0') as usize]); } else { let l = sc.next::<usize>() - 1; let r = sc.next::<usize>(); let ans = seg.find(l, r).3[0][0]; writeln!(out, "{}", ans).ok(); } } } // ---------- begin scannner ---------- #[allow(dead_code)] mod scanner { use std::str::FromStr; pub struct Scanner<'a> { it: std::str::SplitWhitespace<'a>, } impl<'a> Scanner<'a> { pub fn new(s: &'a String) -> Scanner<'a> { Scanner { it: s.split_whitespace(), } } pub fn next<T: FromStr>(&mut self) -> T { self.it.next().unwrap().parse::<T>().ok().unwrap() } pub fn next_bytes(&mut self) -> Vec<u8> { self.it.next().unwrap().bytes().collect() } pub fn next_chars(&mut self) -> Vec<char> { self.it.next().unwrap().chars().collect() } pub fn next_vec<T: FromStr>(&mut self, len: usize) -> Vec<T> { (0..len).map(|_| self.next()).collect() } } } // ---------- end scannner ---------- use std::io::Write; use std::collections::*; type Map<K, V> = BTreeMap<K, V>; type Set<T> = BTreeSet<T>; type Deque<T> = VecDeque<T>; fn main() { use std::io::Read; let mut s = String::new(); std::io::stdin().read_to_string(&mut s).unwrap(); let mut sc = scanner::Scanner::new(&s); let out = std::io::stdout(); let mut out = std::io::BufWriter::new(out.lock()); run(&mut sc, &mut out); } // ---------- begin modint ---------- 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 } #[derive(Clone, Copy, PartialEq, Eq, Default)] pub struct ModInt<const M: u32>(u32); impl<const M: u32> ModInt<{ 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)); const PRIMITIVE_ROOT: u32 = primitive_root(M); const ORDER: usize = 1 << (M - 1).trailing_zeros(); const fn reduce(x: u64) -> u32 { let _ = Self::IS_PRIME; 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 } const fn multiply(a: u32, b: u32) -> u32 { Self::reduce(a as u64 * b as u64) } pub const fn new(v: u32) -> Self { assert!(v < M); Self(Self::reduce(v as u64 * Self::INI)) } pub const fn const_mul(&self, rhs: Self) -> Self { Self(Self::multiply(self.0, rhs.0)) } pub const fn pow(&self, mut n: u64) -> Self { let mut t = Self::new(1); let mut r = *self; while n > 0 { if n & 1 == 1 { t = t.const_mul(r); } r = r.const_mul(r); n >>= 1; } t } pub const fn inv(&self) -> Self { assert!(self.0 != 0); self.pow(M as u64 - 2) } pub const fn get(&self) -> u32 { Self::reduce(self.0 as u64) } pub const fn zero() -> Self { Self::new(0) } pub const fn one() -> Self { Self::new(1) } } impl<const M: u32> Add for ModInt<{ M }> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { let mut v = self.0 + rhs.0; if v >= M { v -= M; } Self(v) } } impl<const M: u32> Sub for ModInt<{ M }> { type Output = Self; fn sub(self, rhs: Self) -> Self::Output { let mut v = self.0 - rhs.0; if self.0 < rhs.0 { v += M; } Self(v) } } impl<const M: u32> Mul for ModInt<{ M }> { type Output = Self; fn mul(self, rhs: Self) -> Self::Output { self.const_mul(rhs) } } impl<const M: u32> Div for ModInt<{ M }> { type Output = Self; fn div(self, rhs: Self) -> Self::Output { self * rhs.inv() } } impl<const M: u32> AddAssign for ModInt<{ M }> { fn add_assign(&mut self, rhs: Self) { *self = *self + rhs; } } impl<const M: u32> SubAssign for ModInt<{ M }> { fn sub_assign(&mut self, rhs: Self) { *self = *self - rhs; } } impl<const M: u32> MulAssign for ModInt<{ M }> { fn mul_assign(&mut self, rhs: Self) { *self = *self * rhs; } } impl<const M: u32> DivAssign for ModInt<{ M }> { fn div_assign(&mut self, rhs: Self) { *self = *self / rhs; } } impl<const M: u32> Neg for ModInt<{ M }> { type Output = Self; fn neg(self) -> Self::Output { if self.0 == 0 { self } else { Self(M - self.0) } } } impl<const M: u32> std::fmt::Display for ModInt<{ M }> { fn fmt<'a>(&self, f: &mut std::fmt::Formatter<'a>) -> std::fmt::Result { write!(f, "{}", self.get()) } } impl<const M: u32> std::fmt::Debug for ModInt<{ M }> { fn fmt<'a>(&self, f: &mut std::fmt::Formatter<'a>) -> std::fmt::Result { write!(f, "{}", self.get()) } } impl<const M: u32> std::str::FromStr for ModInt<{ M }> { type Err = std::num::ParseIntError; fn from_str(s: &str) -> Result<Self, Self::Err> { let val = s.parse::<u32>()?; Ok(ModInt::new(val)) } } impl<const M: u32> From<usize> for ModInt<{ M }> { fn from(val: usize) -> ModInt<{ M }> { ModInt::new((val % M as usize) as u32) } } // ---------- end modint ---------- // ---------- begin precalc ---------- pub struct Precalc<const MOD: u32> { fact: Vec<ModInt<MOD>>, ifact: Vec<ModInt<MOD>>, inv: Vec<ModInt<MOD>>, } impl<const MOD: u32> Precalc<MOD> { pub fn new(size: usize) -> Self { let mut fact = vec![ModInt::one(); size + 1]; let mut ifact = vec![ModInt::one(); size + 1]; let mut inv = vec![ModInt::one(); size + 1]; for i in 2..=size { fact[i] = fact[i - 1] * ModInt::from(i); } ifact[size] = fact[size].inv(); for i in (2..=size).rev() { inv[i] = ifact[i] * fact[i - 1]; ifact[i - 1] = ifact[i] * ModInt::from(i); } Self { fact, ifact, inv } } pub fn fact(&self, n: usize) -> ModInt<MOD> { self.fact[n] } pub fn ifact(&self, n: usize) -> ModInt<MOD> { self.ifact[n] } pub fn inv(&self, n: usize) -> ModInt<MOD> { assert!(0 < n); self.inv[n] } pub fn perm(&self, n: usize, k: usize) -> ModInt<MOD> { if k > n { return ModInt::zero(); } self.fact[n] * self.ifact[n - k] } pub fn binom(&self, n: usize, k: usize) -> ModInt<MOD> { if n < k { return ModInt::zero(); } self.fact[n] * self.ifact[k] * self.ifact[n - k] } } // ---------- end precalc ---------- impl<const M: u32> Zero for ModInt<{ M }> { fn zero() -> Self { Self::zero() } fn is_zero(&self) -> bool { self.0 == 0 } } impl<const M: u32> One for ModInt<{ M }> { fn one() -> Self { Self::one() } fn is_one(&self) -> bool { self.get() == 1 } } // ---------- begin array op ---------- struct NTTPrecalc<const M: u32> { sum_e: [ModInt<{ M }>; 30], sum_ie: [ModInt<{ M }>; 30], } impl<const M: u32> NTTPrecalc<{ M }> { const fn new() -> Self { let cnt2 = (M - 1).trailing_zeros() as usize; let root = ModInt::new(ModInt::<{ M }>::PRIMITIVE_ROOT); let zeta = root.pow((M - 1) as u64 >> cnt2); let mut es = [ModInt::zero(); 30]; let mut ies = [ModInt::zero(); 30]; let mut sum_e = [ModInt::zero(); 30]; let mut sum_ie = [ModInt::zero(); 30]; let mut e = zeta; let mut ie = e.inv(); let mut i = cnt2; while i >= 2 { es[i - 2] = e; ies[i - 2] = ie; e = e.const_mul(e); ie = ie.const_mul(ie); i -= 1; } let mut now = ModInt::one(); let mut inow = ModInt::one(); let mut i = 0; while i < cnt2 - 1 { sum_e[i] = es[i].const_mul(now); sum_ie[i] = ies[i].const_mul(inow); now = ies[i].const_mul(now); inow = es[i].const_mul(inow); i += 1; } Self { sum_e, sum_ie } } } struct NTTPrecalcHelper<const MOD: u32>; impl<const MOD: u32> NTTPrecalcHelper<MOD> { const A: NTTPrecalc<MOD> = NTTPrecalc::new(); } 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 + One + 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 (res, b) in res[i..].iter_mut().zip(rhs.iter()) { *res = *res + *a * *b; } } res } } // transform でlen=1を指定すればNTTになる pub trait ArrayConvolution { type Item; fn transform(&mut self, len: usize); fn inverse_transform(&mut self, len: usize); fn convolution(&self, rhs: &[Self::Item]) -> Vec<Self::Item>; } impl<const M: u32> ArrayConvolution for [ModInt<{ M }>] { type Item = ModInt<{ M }>; fn transform(&mut self, len: usize) { let f = self; let n = f.len(); let k = (n / len).trailing_zeros() as usize; assert!(len << k == n); assert!(k <= ModInt::<{ M }>::ORDER); let pre = &NTTPrecalcHelper::<{ M }>::A; for ph in 1..=k { let p = len << (k - ph); let mut now = ModInt::one(); for (i, f) in f.chunks_exact_mut(2 * p).enumerate() { let (x, y) = f.split_at_mut(p); for (x, y) in x.iter_mut().zip(y.iter_mut()) { let l = *x; let r = *y * now; *x = l + r; *y = l - r; } now *= pre.sum_e[(!i).trailing_zeros() as usize]; } } } fn inverse_transform(&mut self, len: usize) { let f = self; let n = f.len(); let k = (n / len).trailing_zeros() as usize; assert!(len << k == n); assert!(k <= ModInt::<{ M }>::ORDER); let pre = &NTTPrecalcHelper::<{ M }>::A; for ph in (1..=k).rev() { let p = len << (k - ph); let mut inow = ModInt::one(); for (i, f) in f.chunks_exact_mut(2 * p).enumerate() { let (x, y) = f.split_at_mut(p); for (x, y) in x.iter_mut().zip(y.iter_mut()) { let l = *x; let r = *y; *x = l + r; *y = (l - r) * inow; } inow *= pre.sum_ie[(!i).trailing_zeros() as usize]; } } let ik = ModInt::new(2).inv().pow(k as u64); for f in f.iter_mut() { *f *= ik; } } fn convolution(&self, rhs: &[Self::Item]) -> Vec<Self::Item> { if self.len().min(rhs.len()) <= 32 { return self.mul(rhs); } const PARAM: usize = 10; let size = self.len() + rhs.len() - 1; let mut k = 0; while (size + (1 << k) - 1) >> k > PARAM { k += 1; } let len = (size + (1 << k) - 1) >> k; let mut f = vec![ModInt::zero(); len << k]; let mut g = vec![ModInt::zero(); len << k]; f[..self.len()].copy_from_slice(self); g[..rhs.len()].copy_from_slice(rhs); f.transform(len); g.transform(len); let mut buf = [ModInt::zero(); 2 * PARAM - 1]; let buf = &mut buf[..(2 * len - 1)]; let pre = &NTTPrecalcHelper::<{ M }>::A; let mut now = ModInt::one(); for (i, (f, g)) in f .chunks_exact_mut(2 * len) .zip(g.chunks_exact(2 * len)) .enumerate() { let mut r = now; for (f, g) in f.chunks_exact_mut(len).zip(g.chunks_exact(len)) { buf.fill(ModInt::zero()); for (i, f) in f.iter().enumerate() { for (buf, g) in buf[i..].iter_mut().zip(g.iter()) { *buf = *buf + *f * *g; } } f.copy_from_slice(&buf[..len]); for (f, buf) in f.iter_mut().zip(buf[len..].iter()) { *f = *f + r * *buf; } r = -r; } now *= pre.sum_e[(!i).trailing_zeros() as usize]; } f.inverse_transform(len); f.truncate(self.len() + rhs.len() - 1); f } } // ---------- end array op ---------- // ---------- begin trait ---------- 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 Group: Zero + Sub<Output = Self> + Neg<Output = Self> {} pub trait SemiRing: Zero + One {} pub trait Ring: SemiRing + Group {} pub trait Field: Ring + Div<Output = Self> {} impl<T> Group for T where T: Zero + Sub<Output = Self> + Neg<Output = Self> {} impl<T> SemiRing for T where T: Zero + One {} impl<T> Ring for T where T: SemiRing + Group {} impl<T> Field for T where T: Ring + Div<Output = Self> {} pub fn zero<T: Zero>() -> T { T::zero() } pub fn one<T: One>() -> T { T::one() } pub fn pow<T: One + Clone>(mut r: T, mut n: usize) -> T { let mut t = one(); while n > 0 { if n & 1 == 1 { t = t * r.clone(); } r = r.clone() * r; n >>= 1; } t } pub fn pow_sum<T: SemiRing + Clone>(r: T, n: usize) -> T { if n == 0 { T::zero() } else if n & 1 == 1 { T::one() + r.clone() * pow_sum(r, n - 1) } else { let a = T::one() + r.clone(); let b = r.clone() * r; a * pow_sum(b, n / 2) } } // ---------- end trait ---------- // ---------- begin segment tree Point Update Range Query ---------- pub struct SegmentTreePURQ<T, F> { n: usize, size: usize, data: Vec<T>, e: T, op: F, } impl<T, F> SegmentTreePURQ<T, F> where T: Clone, F: Fn(&T, &T) -> T, { pub fn new(n: usize, e: T, op: F) -> Self { assert!(n > 0); let size = n.next_power_of_two(); let data = vec![e.clone(); 2 * size]; SegmentTreePURQ { n, size, data, e, op, } } pub fn update_tmp(&mut self, x: usize, v: T) { assert!(x < self.n); self.data[x + self.size] = v; } pub fn update_all(&mut self) { for i in (1..self.size).rev() { self.data[i] = (self.op)(&self.data[2 * i], &self.data[2 * i + 1]); } } pub fn update(&mut self, x: usize, v: T) { assert!(x < self.n); let mut x = x + self.size; self.data[x] = v; x >>= 1; while x > 0 { self.data[x] = (self.op)(&self.data[2 * x], &self.data[2 * x + 1]); x >>= 1; } } pub fn find(&self, l: usize, r: usize) -> T { assert!(l <= r && r <= self.n); if l == r { return self.e.clone(); } let mut l = self.size + l; let mut r = self.size + r; let mut x = self.e.clone(); let mut y = self.e.clone(); while l < r { if l & 1 == 1 { x = (self.op)(&x, &self.data[l]); l += 1; } if r & 1 == 1 { r -= 1; y = (self.op)(&self.data[r], &y); } l >>= 1; r >>= 1; } (self.op)(&x, &y) } pub fn max_right<P>(&self, l: usize, f: P) -> usize where P: Fn(&T) -> bool, { assert!(l <= self.n); assert!(f(&self.e)); if l == self.n { return self.n; } let mut l = l + self.size; let mut sum = self.e.clone(); while { l >>= l.trailing_zeros(); let v = (self.op)(&sum, &self.data[l]); if !f(&v) { while l < self.size { l <<= 1; let v = (self.op)(&sum, &self.data[l]); if f(&v) { sum = v; l += 1; } } return l - self.size; } sum = v; l += 1; l.count_ones() > 1 } {} self.n } pub fn min_left<P>(&self, r: usize, f: P) -> usize where P: Fn(&T) -> bool, { assert!(r <= self.n); assert!(f(&self.e)); if r == 0 { return 0; } let mut r = r + self.size; let mut sum = self.e.clone(); while { r -= 1; while r > 1 && r & 1 == 1 { r >>= 1; } let v = (self.op)(&self.data[r], &sum); if !f(&v) { while r < self.size { r = 2 * r + 1; let v = (self.op)(&self.data[r], &sum); if f(&v) { sum = v; r -= 1; } } return r + 1 - self.size; } sum = v; (r & (!r + 1)) != r } {} 0 } } // ---------- end segment tree Point Update Range Query ---------- // ---------- begin const matrix ---------- #[derive(Clone, Copy, Debug)] 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) } pub fn swap_row(&mut self, x: usize, y: usize) { assert!(x < R && y < R); self.0.swap(x, y); } pub fn swap_col(&mut self, x: usize, y: usize) { assert!(x < C && y < C); for mat in self.iter_mut() { mat.swap(x, y); } } } impl<T, const R: usize, const C: usize> Matrix<T, R, C> where T: Mul<Output = T> + Copy, { pub fn scalar(&self, k: T) -> Self { let mut res = *self; for a in res.iter_mut().flatten() { *a = *a * k; } res } } 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.iter().flatten().all(|a| a.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::zero(); for (i, a) in res.iter_mut().enumerate() { a[i] = T::one(); } res } fn is_one(&self) -> bool { self.iter().enumerate().all(|(i, a)| { a.iter() .enumerate() .all(|(j, a)| (i == j && a.is_one()) || (i != j && a.is_zero())) }) } } impl<T, const R: usize, const C: usize> AddAssign for Matrix<T, R, C> where T: Add<Output = T> + Copy, { fn add_assign(&mut self, rhs: Self) { for (a, b) in self.iter_mut().zip(rhs.iter()) { for (a, b) in a.iter_mut().zip(b.iter()) { *a = *a + *b; } } } } impl<T, const R: usize, const C: usize> SubAssign for Matrix<T, R, C> where T: Sub<Output = T> + Copy, { fn sub_assign(&mut self, rhs: Self) { for (a, b) in self.iter_mut().zip(rhs.iter()) { for (a, b) in a.iter_mut().zip(b.iter()) { *a = *a - *b; } } } } 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(mut self, rhs: Self) -> Self::Output { self += rhs; self } } 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(mut self, rhs: Self) -> Self::Output { self -= rhs; self } } impl<T, const ROW: usize, const COL: usize> Matrix<T, ROW, COL> where T: Zero + Mul<Output = T> + Copy, { pub fn matmul<const NCOL: usize>(&self, rhs: &Matrix<T, COL, NCOL>) -> Matrix<T, ROW, NCOL> { let mut res = Matrix::<T, ROW, NCOL>::zero(); for (res, a) in res.iter_mut().zip(self.iter()) { for (a, b) in a.iter().zip(rhs.iter()) { for (res, b) in res.iter_mut().zip(b.iter()) { *res = *res + *a * *b; } } } res } } impl<T, const ROW: usize, const COL: usize, const MID: usize> Mul<Matrix<T, MID, COL>> for Matrix<T, ROW, MID> where T: Zero + Mul<Output = T> + Copy, { type Output = Matrix<T, ROW, COL>; fn mul(self, rhs: Matrix<T, MID, COL>) -> Self::Output { self.matmul(&rhs) } } impl<T, const R: usize, const C: usize> MulAssign<Matrix<T, C, C>> for Matrix<T, R, C> where T: Zero + Mul<Output = T> + Copy, { fn mul_assign(&mut self, rhs: Matrix<T, C, C>) { *self = self.matmul(&rhs); } } impl<T, const R: usize, const C: usize> Deref for Matrix<T, R, C> { type Target = [[T; C]; R]; fn deref(&self) -> &Self::Target { &self.0 } } impl<T, const R: usize, const C: usize> DerefMut for Matrix<T, R, C> { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } // ---------- end const matrix ----------