Introduction
Rust has fancy ability to abstract algorithms. Why don't we use it in PS domain?
Many algorithms and data structures are exhausting and boring to implement every time. You can copy ready-made implementations here and focus on the logic.
Matrix
struct Matrix<T>(Vec<T>, usize);
impl<T> std::ops::Index<usize> for Matrix<T> {
type Output = [T];
fn index(&self, index: usize) -> &Self::Output {
&self.0[index * self.1..][..self.1]
}
}
impl<T> std::ops::IndexMut<usize> for Matrix<T> {
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
&mut self.0[index * self.1..][..self.1]
}
}2-SAT
#[derive(Clone, Copy)]
enum Clause {
Pos(usize),
Neg(usize),
}
impl Clause {
fn as_index(&self) -> usize {
match self {
Self::Pos(i) => i << 1 | 1,
Self::Neg(i) => i << 1,
}
}
fn inv(&self) -> Self {
match *self {
Self::Pos(i) => Self::Neg(i),
Self::Neg(i) => Self::Pos(i),
}
}
}
#[derive(Default)]
struct SAT {
head: Vec<u32>,
link: Vec<u32>,
end: Vec<u32>,
}
impl SAT {
fn reserve_clause(&mut self, n: usize) {
self.head.resize(self.head.len() + n * 2, u32::MAX);
}
fn or(&mut self, l: Clause, r: Clause) {
self.connect(l.inv().as_index(), r.as_index());
self.connect(r.inv().as_index(), l.as_index());
}
fn imply(&mut self, l: Clause, r: Clause) {
self.connect(l.as_index(), r.as_index());
self.connect(r.inv().as_index(), l.inv().as_index());
}
fn connect(&mut self, from: usize, to: usize) {
let p = self.head[from];
self.head[from] = self.link.len() as u32;
self.link.push(p);
self.end.push(to as u32);
}
fn try_assign(&self) -> Option<Vec<bool>> {
let n = self.head.len();
let mut c = 0;
let mut s = vec![];
let mut p = vec![];
let mut pre = vec![u32::MAX; n];
let mut comp = vec![u32::MAX; n];
let mut next = 0;
for i in 0..n {
if pre[i] != u32::MAX {
continue;
}
pre[i] = c;
c += 1;
let mut dfs = vec![(i as u32, self.head[i])];
s.push(i as u32);
p.push(i as u32);
while let Some((u, e)) = dfs.last_mut() {
if let Some(&v) = self.end.get(*e as usize) {
if pre[v as usize] == u32::MAX {
pre[v as usize] = c;
c += 1;
s.push(v);
p.push(v);
dfs.push((v, self.head[v as usize]));
continue;
} else if comp[v as usize] == u32::MAX {
while let Some(&w) = p.last() {
if pre[w as usize] <= pre[v as usize] {
break;
}
p.pop();
}
}
*e = self.link[*e as usize];
} else {
if Some(&*u) == p.last() {
p.pop();
while let Some(v) = s.pop() {
if comp[v as usize ^ 1] == next {
return None;
}
comp[v as usize] = next;
if v == *u {
break;
}
}
next += 1;
}
dfs.pop();
if let Some((_, e)) = dfs.last_mut() {
*e = self.link[*e as usize];
}
}
}
}
Some(comp.chunks(2).map(|c| c[0] > c[1]).collect())
}
}Pseudo Random Number Generator
Reference: https://github.com/BrutPitt/fastPRNG
struct Rng32([u32; 4]);
impl Rng32 {
fn split_mix(v: u32) -> u32 {
let mut z = v.wrapping_add(0x9e3779b9);
z = (z ^ (z >> 15)).wrapping_mul(0x85ebca6b);
z = (z ^ (z >> 13)).wrapping_mul(0xc2b2ae35);
z ^ (z >> 16)
}
fn new() -> Self {
let mut seed = 0;
unsafe { std::arch::x86_64::_rdrand32_step(&mut seed) };
let mut prev = seed;
Self(std::array::from_fn(|_| {
prev = Self::split_mix(prev);
prev
}))
}
fn next(&mut self, n: u32) -> u32 {
let [x, y, z, w] = &mut self.0;
let res = x.wrapping_add(*w);
let t = x.wrapping_shl(9);
*y ^= *x;
*w ^= *y;
*y ^= *z;
*x ^= *w;
*z ^= t;
*w = w.rotate_left(11);
((res as u64 * n as u64) >> 32) as u32
}
}
struct Rng64([u64; 4]);
impl Rng64 {
fn split_mix(v: u64) -> u64 {
let mut z = v.wrapping_add(0x9e3779b97f4a7c15);
z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb);
z ^ (z >> 31)
}
fn new() -> Self {
let mut seed = 0;
unsafe { std::arch::x86_64::_rdrand64_step(&mut seed) };
let mut prev = seed;
Self(std::array::from_fn(|_| {
prev = Self::split_mix(prev);
prev
}))
}
fn next(&mut self, n: u64) -> u64 {
let [x, y, z, w] = &mut self.0;
let res = x.wrapping_add(*w);
let t = x.wrapping_shl(17);
*y ^= *x;
*w ^= *y;
*y ^= *z;
*x ^= *w;
*z ^= t;
*w = w.rotate_left(45);
((res as u128 * n as u128) >> 64) as u64
}
}Arbitrary Precision Integer
#[derive(Clone, PartialEq, Eq)]
struct BigInt(bool, Vec<u64>);
impl std::str::FromStr for BigInt {
type Err = std::convert::Infallible;
fn from_str(mut s: &str) -> std::result::Result<Self, Self::Err> {
let sign = s.starts_with("-");
if sign {
s = &s[1..];
}
let size = (s.len() + RADIX_POW - 1) / RADIX_POW;
let mut v = Vec::with_capacity(size);
for c in s.as_bytes().rchunks(RADIX_POW) {
let mut t = 0;
for &b in c {
t *= 10;
t += (b - b'0') as u64;
}
v.push(t);
}
Ok(Self(sign, v))
}
}
impl std::fmt::Display for BigInt {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if self.0 {
write!(f, "-")?;
}
let mut a = self.1.iter().rev();
let last = *a.next().unwrap();
write!(f, "{last}")?;
for l in a {
write!(f, "{l:00$}", RADIX_POW)?;
}
Ok(())
}
}
impl std::cmp::PartialOrd for BigInt {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl std::cmp::Ord for BigInt {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
use std::cmp::Ordering::*;
if self.0 && !other.0 {
Less
} else if !self.0 && other.0 {
Greater
} else {
let ord = if self.1.len() != other.1.len() {
self.1.len().cmp(&other.1.len())
} else {
let mut last = Equal;
for (l, r) in self.1.iter().rev().zip(other.1.iter().rev()) {
last = l.cmp(r);
if last != Equal {
break;
}
}
last
};
if self.0 {
match ord {
Less => Greater,
Equal => Equal,
Greater => Less,
}
} else {
ord
}
}
}
}
const RADIX: u64 = 1_000_000;
const RADIX_POW: usize = 6;
impl BigInt {
fn new() -> Self {
Self(false, vec![0])
}
fn is_zero(&self) -> bool {
self.1.len() == 1 && self.1[0] == 0
}
fn normalize(&mut self) {
let len = self.1.iter().rposition(|&l| l != 0).unwrap_or(0) + 1;
self.1.truncate(len);
if self.1.len() == 1 && self.1[0] == 0 {
self.0 = false;
}
}
fn add(&mut self, rhs: &Self) {
if self.0 == rhs.0 {
self.unsigned_add_assign(rhs);
} else {
self.unsigned_sub_assign(rhs);
}
}
fn sub(&mut self, rhs: &Self) {
if self.0 == rhs.0 {
self.unsigned_sub_assign(rhs);
} else {
self.unsigned_add_assign(rhs);
}
}
fn divmod_scalar(&mut self, rhs: u64) -> u64 {
let mut borrow = 0;
for d in self.1.iter_mut().rev() {
borrow *= RADIX;
borrow += *d;
(*d, borrow) = (borrow / rhs, borrow % rhs);
}
self.normalize();
borrow
}
fn mul(&mut self, rhs: &Self) {
const NTT_0_INV: u64 = 285212675;
let pad = (self.1.len() + rhs.1.len() - 1).next_power_of_two();
self.1.resize(pad, 0);
let mut l = &mut self.1;
let mut r = vec![0; pad];
r[..rhs.1.len()].copy_from_slice(&rhs.1);
let mut l1 = l.clone();
let mut r1 = r.clone();
NTT[0].run(&mut l, false);
NTT[0].run(&mut r, false);
for (l, r) in l.iter_mut().zip(r) {
*l = NTT[0].m.rem(*l as u64 * r as u64);
}
NTT[0].run(&mut l, true);
NTT[1].run(&mut l1, false);
NTT[1].run(&mut r1, false);
for (l, r) in l1.iter_mut().zip(r1) {
*l = NTT[1].m.rem(*l as u64 * r as u64);
}
NTT[1].run(&mut l1, true);
for (l, l1) in l.iter_mut().zip(l1) {
let r = NTT[1].m.rem((l1 + NTT[1].m.1 - *l) * NTT_0_INV);
*l += NTT[0].m.1 * r;
}
let mut carry = 0;
for x in l {
*x += carry;
carry = *x / RADIX;
*x %= RADIX;
}
while carry > 0 {
self.1.push(carry % RADIX);
carry /= RADIX;
}
self.0 = self.0 != rhs.0;
self.normalize();
}
fn unsigned_add_assign(&mut self, rhs: &Self) {
if rhs.is_zero() {
return;
} else if self.is_zero() {
self.0 = rhs.0;
self.1 = rhs.1.clone();
return;
}
let n = self.1.len().min(rhs.1.len());
let mut carry = 0;
for (l, &r) in self.1[..n].iter_mut().zip(&rhs.1) {
*l += r;
*l += carry;
if *l >= RADIX {
*l -= RADIX;
carry = 1;
} else {
carry = 0;
}
}
for l in &mut self.1[n..] {
*l += carry;
if *l >= RADIX {
*l -= RADIX;
carry = 1;
} else {
carry = 0;
break;
}
}
for &r in &rhs.1[n..] {
let mut l = r + carry;
if l >= RADIX {
l -= RADIX;
carry = 1;
} else {
carry = 0;
}
self.1.push(l);
}
if carry > 0 {
self.1.push(1);
}
self.normalize();
}
fn unsigned_sub_assign(&mut self, rhs: &Self) {
if rhs.is_zero() {
return;
} else if self.is_zero() {
self.0 = !rhs.0;
self.1 = rhs.1.clone();
return;
}
let less = self.1.len().saturating_sub(rhs.1.len());
let mut a = self.1.iter_mut();
let mut b = rhs.1.iter().chain(std::iter::repeat(&0).take(less));
let mut zero = true;
let mut carry = 0;
for (l, &r) in a.by_ref().zip(b.by_ref()) {
if zero {
if r != 0 {
zero = false;
*l += RADIX - r;
} else {
continue;
}
} else {
*l += RADIX - 1 + carry - r;
}
if *l >= RADIX {
*l -= RADIX;
carry = 1;
} else {
carry = 0;
}
}
for &r in b {
let mut add;
if zero {
if r != 0 {
zero = false;
add = RADIX + carry - r;
} else {
continue;
}
} else {
add = RADIX - 1 + carry - r;
}
if add >= RADIX {
add -= RADIX;
carry = 1;
} else {
carry = 0;
}
self.1.push(add);
}
if carry == 0 {
self.0 = !self.0;
let mut a = self.1.iter_mut().skip_while(|v| **v == 0);
if let Some(first) = a.next() {
*first = RADIX - *first;
for v in a {
*v = RADIX - 1 - *v;
}
}
}
self.normalize();
}
}
struct NttParam {
m: ModU64,
u: u64,
ui: u64,
s: u32,
}
static NTT: &[NttParam] = &[
NttParam {
m: ModU64::new(1107296257),
u: 1087287097,
ui: 623044540,
s: 25,
},
NttParam {
m: ModU64::new(1711276033),
u: 969788637,
ui: 1790856,
s: 25,
},
];
impl NttParam {
fn run(&self, a: &mut [u64], inv: bool) {
let n = a.len();
if n == 1 {
return;
}
let s = n.leading_zeros() + 1;
for i in 0..n {
let r = i.reverse_bits() >> s;
if i < r {
a.swap(i, r);
}
}
let u = if inv { self.ui } else { self.u };
for k in 1..=n.trailing_zeros() {
let mut wlen = u;
for _ in k..self.s {
wlen = self.m.rem(wlen * wlen);
}
let kh = 1 << (k - 1);
for i in 0..(n + (1 << k) - 1) >> k {
let i = i << k;
let mut w = 1;
for j in 0..kh {
let u = a[i + j];
let v = self.m.rem(a[i + j + kh] * w);
let mut s = u + v;
if s >= self.m.1 {
s -= self.m.1;
}
a[i + j] = s;
let mut d = u + self.m.1 - v;
if d >= self.m.1 {
d -= self.m.1;
}
a[i + j + kh] = d;
w = self.m.rem(w * wlen);
}
}
}
if inv {
let p = self.m.1 as i64;
let ni = ((egcd(n as i64, p).1 % p + p) % p) as u64;
for x in a {
*x = self.m.rem(*x as u64 * ni);
}
}
}
}
fn egcd(mut a: i64, mut b: i64) -> (i64, i64, i64) {
let (mut x, mut y, mut x1, mut y1) = (1, 0, 0, 1);
while b != 0 {
let q = a / b;
(x, x1) = (x1, x - q * x1);
(y, y1) = (y1, y - q * y1);
(a, b) = (b, a - q * b);
}
(a, x, y)
}
#[derive(Copy, Clone)]
struct ModU64(u128, u64);
impl ModU64 {
const fn new(div: u64) -> Self {
Self((!0u128 / div as u128).wrapping_add(1), div)
}
fn multop(a: u128, b: u64) -> u64 {
let mut bottom = (a as u64 as u128) * b as u128;
bottom >>= 64;
let top = (a >> 64) * b as u128;
((bottom + top) >> 64) as u64
}
fn rem(&self, a: u64) -> u64 {
let low = self.0.wrapping_mul(a as u128);
Self::multop(low, self.1)
}
}Bitset
struct Bitset(Vec<std::arch::x86_64::__m256i>, usize);
impl Bitset {
#[target_feature(enable = "avx2")]
unsafe fn new(n: usize) -> Self {
use std::arch::x86_64::*;
let mut v = vec![];
v.resize_with((n + 255) / 256, || unsafe { _mm256_setzero_si256() });
Self(v, n)
}
fn len(&self) -> usize {
self.1
}
fn get(&self, i: usize) -> bool {
let b64 =
unsafe { std::slice::from_raw_parts(self.0.as_ptr() as *const u64, self.0.len() * 4) };
b64[i / 64] & (1 << (i % 64)) != 0
}
fn set(&mut self, i: usize, v: bool) {
let b64 = unsafe {
std::slice::from_raw_parts_mut(self.0.as_mut_ptr() as *mut u64, self.0.len() * 4)
};
if v {
b64[i / 64] |= 1 << (i % 64);
} else {
b64[i / 64] &= !(1 << (i % 64));
}
}
fn flip(&mut self, i: usize) {
let b64 = unsafe {
std::slice::from_raw_parts_mut(self.0.as_mut_ptr() as *mut u64, self.0.len() * 4)
};
b64[i / 64] ^= 1 << (i % 64);
}
fn shl(&self, x: usize) -> Self {
let mut shl = unsafe { Self::new(self.len() + x) };
let shl_b64 = unsafe {
std::slice::from_raw_parts_mut(shl.0.as_mut_ptr().cast::<u64>(), (shl.len() + 63) / 64)
};
let b64 = unsafe {
std::slice::from_raw_parts(self.0.as_ptr().cast::<u64>(), (self.len() + 63) / 64)
};
let x_chunk = x / 64;
let x_inside = (x % 64) as u32;
if x_inside == 0 {
for (dest, &src) in shl_b64.iter_mut().skip(x_chunk).zip(b64) {
*dest = src.wrapping_shl(x_inside);
}
} else {
let mut low = 0;
for (dest, &src) in shl_b64.iter_mut().skip(x_chunk).zip(b64) {
*dest = src.wrapping_shl(x_inside) | low;
low = src >> (64 - x_inside);
}
if shl_b64.len() > b64.len() + x_chunk {
*shl_b64.last_mut().unwrap() = low;
}
}
shl
}
#[target_feature(enable = "avx2")]
unsafe fn or(&mut self, other: &Self) {
for (a, &b) in self.0.iter_mut().zip(&other.0) {
*a = std::arch::x86_64::_mm256_or_si256(*a, b);
}
}
fn split_range(mut begin: usize, end: usize) -> Option<(usize, usize)> {
let back = end & !255;
if begin & 255 != 0 {
begin += 256 - (begin & 255);
}
if begin < back {
Some((begin, back))
} else {
None
}
}
#[target_feature(enable = "avx2")]
unsafe fn find_first_unset(&self, l: usize, r: usize) -> Option<usize> {
let b64 =
unsafe { std::slice::from_raw_parts(self.0.as_ptr() as *const u64, self.0.len() * 4) };
if let Some((le, rb)) = Self::split_range(l, r) {
for i in l..le {
if b64[i >> 6] & (1 << (i & 63)) == 0 {
return Some(i);
}
}
let mut mask = 0;
if let Some(mut i) = self.0[le >> 8..rb >> 8].iter().position(|p| {
use std::arch::x86_64::*;
let v256 = unsafe { _mm256_load_si256(p) };
let zero = unsafe { _mm256_set1_epi8(-1) };
let nonzero = unsafe { _mm256_cmpeq_epi8(zero, v256) };
mask = unsafe { _mm256_movemask_epi8(nonzero) } as u32;
mask != !0
}) {
i += le >> 8;
let j = (0..32).position(|s| mask & (1 << s) == 0).unwrap();
let b8 = unsafe {
std::slice::from_raw_parts(self.0.as_ptr() as *const u8, self.0.len() * 32)
};
let k = b8[i * 32 + j].trailing_ones() as usize;
let begin = (i * 32 + j) * 8 + k;
return Some(begin);
}
for i in rb..r {
if b64[i >> 6] & (1 << (i & 63)) == 0 {
return Some(i);
}
}
None
} else {
for i in l..r {
if b64[i >> 6] & (1 << (i & 63)) == 0 {
return Some(i);
}
}
None
}
}
#[target_feature(enable = "avx2")]
unsafe fn find_first_set(&self, l: usize, r: usize) -> Option<usize> {
let b64 =
unsafe { std::slice::from_raw_parts(self.0.as_ptr() as *const u64, self.0.len() * 4) };
if let Some((le, rb)) = Self::split_range(l, r) {
for i in l..le {
if b64[i >> 6] & (1 << (i & 63)) != 0 {
return Some(i);
}
}
let mut mask = 0;
if let Some(mut i) = self.0[le >> 8..rb >> 8].iter().position(|p| {
use std::arch::x86_64::*;
let v256 = unsafe { _mm256_load_si256(p) };
let zero = unsafe { _mm256_setzero_si256() };
let nonzero = unsafe { _mm256_cmpeq_epi8(zero, v256) };
mask = unsafe { _mm256_movemask_epi8(nonzero) } as u32;
mask != !0
}) {
i += le >> 8;
let j = (0..32).position(|s| mask & (1 << s) == 0).unwrap();
let b8 = unsafe {
std::slice::from_raw_parts(self.0.as_ptr() as *const u8, self.0.len() * 32)
};
let k = b8[i * 32 + j].trailing_zeros() as usize;
let begin = (i * 32 + j) * 8 + k;
return Some(begin);
}
for i in rb..r {
if b64[i >> 6] & (1 << (i & 63)) != 0 {
return Some(i);
}
}
None
} else {
for i in l..r {
if b64[i >> 6] & (1 << (i & 63)) != 0 {
return Some(i);
}
}
None
}
}
#[target_feature(enable = "avx2")]
unsafe fn flip_range(&mut self, l: usize, r: usize) {
let b64 = std::slice::from_raw_parts_mut(self.0.as_mut_ptr() as *mut u64, self.0.len() * 4);
if let Some((le, rb)) = Self::split_range(l, r) {
use std::arch::x86_64::*;
for i in l..le {
b64[i >> 6] ^= 1 << (i & 63);
}
let one = _mm256_set1_epi8(-1);
for v in &mut self.0[le >> 8..rb >> 8] {
let load = _mm256_load_si256(v);
let xor = _mm256_xor_si256(load, one);
_mm256_store_si256(v, xor);
}
for i in rb..r {
b64[i >> 6] ^= 1 << (i & 63);
}
} else {
for i in l..r {
b64[i >> 6] ^= 1 << (i & 63);
}
}
}
#[target_feature(enable = "avx2")]
unsafe fn count_ones(&mut self, l: usize, r: usize) -> u32 {
let b64 = std::slice::from_raw_parts_mut(self.0.as_mut_ptr() as *mut u64, self.0.len() * 4);
let mut count = 0;
if let Some((le, rb)) = Self::split_range(l, r) {
use std::arch::x86_64::*;
for i in l..le {
if b64[i >> 6] & 1 << (i & 63) != 0 {
count += 1;
}
}
let b0 = _mm256_set1_epi32(0x55555555);
let b1 = _mm256_set1_epi32(0x33333333);
let b2 = _mm256_set1_epi32(0x0f0f0f0f);
let b3 = _mm256_set1_epi32(0x01010101);
let mut vcount = _mm256_setzero_si256();
for v in &mut self.0[le >> 8..rb >> 8] {
let load = _mm256_load_si256(v);
let c1 = _mm256_sub_epi32(load, _mm256_and_si256(b0, _mm256_srli_epi32::<1>(load)));
let c2 = _mm256_add_epi32(
_mm256_and_si256(c1, b1),
_mm256_and_si256(_mm256_srli_epi32::<2>(c1), b1),
);
let c3 = _mm256_srli_epi32::<24>(_mm256_mullo_epi32(
_mm256_and_si256(_mm256_add_epi32(c2, _mm256_srli_epi32::<4>(c2)), b2),
b3,
));
vcount = _mm256_add_epi32(vcount, c3);
}
let [a, b, c, d, e, f, g, h]: [u32; 8] = std::mem::transmute(vcount);
count += a + b + c + d + e + f + g + h;
for i in rb..r {
if b64[i >> 6] & 1 << (i & 63) != 0 {
count += 1;
}
}
} else {
for i in l..r {
if b64[i >> 6] & 1 << (i & 63) != 0 {
count += 1;
}
}
}
count
}
}
impl std::fmt::Debug for Bitset {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let b64 =
unsafe { std::slice::from_raw_parts(self.0.as_ptr().cast::<u64>(), self.0.len() * 4) };
let mut iter = b64.iter().rev().skip_while(|&&b| b == 0);
if let Some(&first) = iter.next() {
write!(f, "{first:b}")?;
for &b in iter {
write!(f, "{b:064b}")?;
}
Ok(())
} else {
write!(f, "0")
}
}
}Fast IO
struct Reader {
begin: *const u8,
cur: *const u8,
end: *const u8,
goff: usize,
cap: usize,
}
impl Reader {
fn new(capacity: usize) -> Self {
let begin = unsafe { mmap(0, capacity, 3, 2, 0, 0) };
Self {
begin,
end: unsafe { begin.add(capacity) },
cur: begin,
goff: 0,
cap: capacity,
}
}
fn try_refill(&mut self, readahead: usize) {
if unsafe { self.cur.add(readahead) } < self.end {
return;
}
self.goff += unsafe { self.cur.offset_from(self.begin) } as usize;
let add = self.goff & 4095;
self.goff &= !4095;
self.begin = unsafe { mmap(self.begin as usize, self.cap, 3, 18, 0, self.goff as isize) };
self.cur = unsafe { self.begin.add(add) };
self.end = unsafe { self.begin.add(self.cap) };
}
fn u32(&mut self) -> u32 {
let mut c = unsafe { self.cur.cast::<u64>().read_unaligned() };
let m = !c & 0x1010101010101010;
let len = m.trailing_zeros() >> 3;
c <<= (8 - len) << 3;
c = (c & 0x0F0F0F0F0F0F0F0F).wrapping_mul(2561) >> 8;
c = (c & 0x00FF00FF00FF00FF).wrapping_mul(6553601) >> 16;
c = (c & 0x0000FFFF0000FFFF).wrapping_mul(42949672960001) >> 32;
self.cur = unsafe { self.cur.add(len as usize) };
if len == 8 {
if unsafe { self.cur.read() } & 0x10 != 0 {
c *= 10;
c += (unsafe { self.cur.read() } - b'0') as u64;
self.cur = unsafe { self.cur.add(1) };
}
if unsafe { self.cur.read() } & 0x10 != 0 {
c *= 10;
c += (unsafe { self.cur.read() } - b'0') as u64;
self.cur = unsafe { self.cur.add(1) };
}
}
self.cur = unsafe { self.cur.add(1) };
c as u32
}
fn i32(&mut self) -> i32 {
let sign = unsafe { self.cur.read() } == b'-';
(if sign {
self.cur = unsafe { self.cur.add(1) };
self.u32().wrapping_neg()
} else {
self.u32()
}) as i32
}
fn u64(&mut self) -> u64 {
let mut c = unsafe { self.cur.cast::<u64>().read_unaligned() };
let m = !c & 0x1010101010101010;
let len = m.trailing_zeros() >> 3;
c <<= (8 - len) << 3;
c = (c & 0x0F0F0F0F0F0F0F0F).wrapping_mul(2561) >> 8;
c = (c & 0x00FF00FF00FF00FF).wrapping_mul(6553601) >> 16;
c = (c & 0x0000FFFF0000FFFF).wrapping_mul(42949672960001) >> 32;
self.cur = unsafe { self.cur.add(len as usize) };
if len == 8 && unsafe { self.cur.read() } & 16 != 0 {
let mut d = unsafe { self.cur.cast::<u64>().read_unaligned() };
let m = !d & 0x1010101010101010;
let len = m.trailing_zeros() >> 3;
for _ in 0..len {
c *= 10;
}
d <<= (8 - len) << 3;
d = (d & 0x0F0F0F0F0F0F0F0F).wrapping_mul(2561) >> 8;
d = (d & 0x00FF00FF00FF00FF).wrapping_mul(6553601) >> 16;
d = (d & 0x0000FFFF0000FFFF).wrapping_mul(42949672960001) >> 32;
c += d;
self.cur = unsafe { self.cur.add(len as usize) };
if len == 8 {
if unsafe { self.cur.read() } & 0x10 != 0 {
c *= 10;
c += (unsafe { self.cur.read() } - b'0') as u64;
self.cur = unsafe { self.cur.add(1) };
}
if unsafe { self.cur.read() } & 0x10 != 0 {
c *= 10;
c += (unsafe { self.cur.read() } - b'0') as u64;
self.cur = unsafe { self.cur.add(1) };
}
if unsafe { self.cur.read() } & 0x10 != 0 {
c *= 10;
c += (unsafe { self.cur.read() } - b'0') as u64;
self.cur = unsafe { self.cur.add(1) };
}
if unsafe { self.cur.read() } & 0x10 != 0 {
c *= 10;
c += (unsafe { self.cur.read() } - b'0') as u64;
self.cur = unsafe { self.cur.add(1) };
}
}
}
self.cur = unsafe { self.cur.add(1) };
c
}
fn i64(&mut self) -> i64 {
let sign = unsafe { self.cur.read() } == b'-';
(if sign {
self.cur = unsafe { self.cur.add(1) };
self.u64().wrapping_neg()
} else {
self.u64()
}) as i64
}
fn consume(&mut self, add: usize) {
self.cur = unsafe { self.cur.add(add) };
}
fn until(&mut self, delim: u8, buf: &mut String) -> usize {
#[target_feature(enable = "avx2,sse4.2")]
unsafe fn memchr(s: &[u8], delim: u8) -> Option<usize> {
s.iter().position(|&b| b == delim)
}
let mut total = 0;
loop {
let len = unsafe { self.end.offset_from(self.cur) } as usize;
let range = unsafe { std::slice::from_raw_parts(self.cur, len) };
if let Some(i) = unsafe { memchr(range, delim) } {
unsafe { buf.as_mut_vec() }.extend_from_slice(&range[..i]);
self.cur = unsafe { self.cur.add(i + 1) };
break total + i;
} else {
unsafe { buf.as_mut_vec() }.extend_from_slice(&range);
self.cur = self.end;
self.try_refill(1);
total += len;
}
}
}
fn remain(&self) -> &[u8] {
let len = unsafe { self.end.offset_from(self.cur) } as usize;
unsafe { std::slice::from_raw_parts(self.cur, len) }
}
fn discard(&mut self, until: u8) -> usize {
let mut len = 0;
#[target_feature(enable = "avx2")]
unsafe fn index(s: &[u8], b: u8) -> Option<usize> {
s.iter().position(|&c| c == b)
}
loop {
let pos = unsafe { index(self.remain(), until) };
if let Some(pos) = pos {
len += pos;
self.cur = unsafe { self.cur.add(pos + 1) };
break len;
}
len += unsafe { self.end.offset_from(self.cur) } as usize;
self.cur = self.end;
self.try_refill(1);
}
}
}
struct Writer {
buf: Vec<u8>,
off: usize,
}
impl Drop for Writer {
fn drop(&mut self) {
self.flush();
}
}
#[repr(align(16))]
struct B128([u8; 16]);
#[target_feature(enable = "avx2")]
unsafe fn cvt8(out: &mut B128, n: u32) -> usize {
use std::arch::x86_64::*;
let x = _mm_cvtsi32_si128(n as i32);
let div_10000 = _mm_set1_epi32(0xd1b71759u32 as i32);
let mul_10000_merge = _mm_set1_epi32(55536);
let div_var = _mm_setr_epi16(
8389,
5243,
13108,
0x8000u16 as i16,
8389,
5243,
13108,
0x8000u16 as i16,
);
let shift_var = _mm_setr_epi16(
1 << 7,
1 << 11,
1 << 13,
(1 << 15) as i16,
1 << 7,
1 << 11,
1 << 13,
(1 << 15) as i16,
);
let mul_10 = _mm_set1_epi16(10);
let ascii0 = _mm_set1_epi8(48);
let x_div_10000 = _mm_srli_epi64::<45>(_mm_mul_epu32(x, div_10000));
let y = _mm_add_epi32(x, _mm_mul_epu32(x_div_10000, mul_10000_merge));
let t0 = _mm_slli_epi16::<2>(_mm_shuffle_epi32::<5>(_mm_unpacklo_epi16(y, y)));
let t1 = _mm_mulhi_epu16(t0, div_var);
let t2 = _mm_mulhi_epu16(t1, shift_var);
let t3 = _mm_slli_epi64::<16>(t2);
let t4 = _mm_mullo_epi16(t3, mul_10);
let t5 = _mm_sub_epi16(t2, t4);
let t6 = _mm_packus_epi16(_mm_setzero_si128(), t5);
let mask = _mm_movemask_epi8(_mm_cmpeq_epi8(t6, _mm_setzero_si128()));
let offset = (mask & !0x8000).trailing_ones() as usize;
let ascii = _mm_add_epi8(t6, ascii0);
_mm_store_si128(out.0.as_mut_ptr().cast(), ascii);
offset
}
impl Writer {
fn new(capacity: usize) -> Self {
let mut buf = vec![];
buf.reserve_exact(capacity);
unsafe { buf.set_len(capacity) };
Self { buf, off: 0 }
}
fn flush(&mut self) {
unsafe { write(1, self.buf.as_ptr(), self.off) };
self.off = 0;
}
fn try_flush(&mut self, readahead: usize) {
if self.off + readahead > self.buf.len() {
self.flush();
}
}
fn byte(&mut self, b: u8) {
self.try_flush(1);
self.buf[self.off] = b;
self.off += 1;
}
fn bytes(&mut self, s: &[u8]) {
let mut i = 0;
while i < s.len() {
let rem = s[i..].len().min(self.buf[self.off..].len());
self.buf[self.off..self.off + rem].copy_from_slice(&s[i..i + rem]);
self.off += rem;
i += rem;
if self.off == self.buf.len() {
self.flush();
}
}
}
fn i32(&mut self, n: i32) {
if n < 0 {
self.byte(b'-');
self.u32((n as u32).wrapping_neg());
} else {
self.u32(n as u32);
}
}
fn u32(&mut self, n: u32) {
let mut b128 = B128([0u8; 16]);
let mut off;
if n >= 100_000_000 {
self.try_flush(10);
let mut hi = n / 100_000_000;
let lo = n % 100_000_000;
unsafe { cvt8(&mut b128, lo) };
off = 8;
off -= 1;
b128.0[off] = (hi % 10) as u8 + b'0';
if hi >= 10 {
off -= 1;
hi /= 10;
b128.0[off] = hi as u8 + b'0';
}
} else {
self.try_flush(8);
off = unsafe { cvt8(&mut b128, n) };
}
let len = 16 - off;
self.buf[self.off..self.off + len].copy_from_slice(&b128.0[off..]);
self.off += len;
}
fn i64(&mut self, n: i64) {
if n < 0 {
self.byte(b'-');
self.u64((n as u64).wrapping_neg());
} else {
self.u64(n as u64);
}
}
fn u64(&mut self, n: u64) {
let mut hi128 = B128([0u8; 16]);
let mut lo128 = B128([0u8; 16]);
let mut hioff;
let looff;
if n >= 10_000_000_000_000_000 {
self.try_flush(20);
let mut hi = (n / 10_000_000_000_000_000) as u32;
let lo = n % 10_000_000_000_000_000;
let lohi = (lo / 100_000_000) as u32;
let lolo = (lo % 100_000_000) as u32;
unsafe { cvt8(&mut hi128, lohi) };
unsafe { cvt8(&mut lo128, lolo) };
hioff = 8;
looff = 8;
hioff -= 1;
hi128.0[hioff] = (hi % 10) as u8 + b'0';
if hi >= 10 {
hioff -= 1;
hi /= 10;
hi128.0[hioff] = (hi % 10) as u8 + b'0';
}
if hi >= 10 {
hioff -= 1;
hi /= 10;
hi128.0[hioff] = (hi % 10) as u8 + b'0';
}
if hi >= 10 {
hioff -= 1;
hi /= 10;
hi128.0[hioff] = hi as u8 + b'0';
}
} else if n >= 100_000_000 {
self.try_flush(16);
let hi = (n / 100_000_000) as u32;
let lo = (n % 100_000_000) as u32;
hioff = unsafe { cvt8(&mut hi128, hi) };
unsafe { cvt8(&mut lo128, lo) };
looff = 8;
} else {
self.try_flush(8);
hioff = 16;
looff = unsafe { cvt8(&mut lo128, n as u32) };
}
let len = 16 - hioff;
self.buf[self.off..self.off + len].copy_from_slice(&hi128.0[hioff..]);
self.off += len;
let len = 16 - looff;
self.buf[self.off..self.off + len].copy_from_slice(&lo128.0[looff..]);
self.off += len;
}
}Scanf
You can use scanf. Not the best API, but I think it fits best for PS area.
fn solve() {
scanf!("%d%d", a: i32, b: i32);
println!("{}", a + b);
scanf!("%s", name: [u8; 32]);
println!("{name}");
}
fn main() {
unsafe {
*stdin |= 0x8000;
STDOUT = Some(BufWriter::with_capacity(1 << 17, File::from_raw_fd(1)));
}
solve();
unsafe {
STDOUT.as_mut().unwrap_unchecked().flush().ok();
_exit(0);
}
}
use std::{fs::File, io::*, os::unix::io::FromRawFd};
static mut STDOUT: Option<BufWriter<File>> = None;
#[macro_export]
macro_rules! println {
($($t:tt)*) => { unsafe { writeln!(STDOUT.as_mut().unwrap_unchecked(), $($t)*).unwrap_unchecked() } };
}
#[macro_export]
macro_rules! print {
($($t:tt)*) => { unsafe { write!(STDOUT.as_mut().unwrap_unchecked(), $($t)*).unwrap_unchecked() } };
}
#[macro_export]
macro_rules! flush {
() => {
unsafe {
STDOUT
.as_mut()
.unwrap_unchecked()
.flush()
.unwrap_unchecked()
}
};
}
#[macro_export]
macro_rules! scanf {
($fmt:literal $(, $($t:tt)+)?) => {
scanf!(@def $($($t)+)?);
scanf!(@call $($($t)+)?, $fmt);
scanf!(@bind $($($t)+)?);
};
(@def) => {};
(@def $(mut)? $name:ident: [u8; $size:expr] $(, $($t:tt)+)?) => {
let mut $name = vec![0u8; $size + 1];
scanf!(@def $($($t)+)?);
};
(@def $(mut)? $name:ident: $ty:ty $(, $($t:tt)+)?) => {
let mut $name = std::mem::MaybeUninit::<$ty>::uninit();
scanf!(@def $($($t)+)?);
};
($($names:ident),* @call $fmt:literal) => {
unsafe {
fscanf(stdin, concat!($fmt, "\0").as_ptr(), $($names.as_mut_ptr()),*);
};
};
($($names:ident),* @call $(mut)? $name:ident: [u8; $size:expr] $(, $($t:tt)+)?) => {
scanf!($($names,)* $name @call $($($t)+)?);
};
($($names:ident),* @call $(mut)? $name:ident: $ty:ty $(, $($t:tt)+)?) => {
scanf!($($names,)* $name @call $($($t)+)?);
};
(@bind) => {};
(@bind $name:ident: [u8; $size:expr] $(, $($t:tt)+)?) => {
$name.pop();
let $name = unsafe { String::from_utf8_unchecked($name) };
scanf!(@bind $($($t)+)?);
};
(@bind mut $name:ident: [u8; $size:expr] $(, $($t:tt)+)?) => {
$name.pop();
let mut $name = unsafe { String::from_utf8_unchecked($name) };
scanf!(@bind $($($t)+)?);
};
(@bind $name:ident: $ty:ty $(, $($t:tt)+)?) => {
let $name = unsafe { $name.assume_init() };
scanf!(@bind $($($t)+)?);
};
(@bind mut $name:ident: $ty:ty $(, $($t:tt)+)?) => {
let mut $name = unsafe { $name.assume_init() };
scanf!(@bind $($($t)+)?);
};
}
#[link(name = "c")]
extern "C" {
fn fscanf(file: *mut usize, fmt: *const u8, ...) -> i32;
fn _exit(code: i32) -> !;
static stdin: *mut usize;
}
Float128
#[derive(Clone, Copy)]
#[allow(non_camel_case_types)]
struct f128(std::arch::x86_64::__m128);
macro_rules! impl_float {
(from $t:ty: $from:ident $to:ident) => {
impl From<$t> for f128 {
fn from(v: $t) -> Self {
unsafe { $from(v) }
}
}
impl From<f128> for $t {
fn from(v: f128) -> Self {
unsafe { $to(v) }
}
}
};
(op $trait:ident $method:ident $func:ident) => {
impl std::ops::$trait<Self> for f128 {
type Output = Self;
fn $method(self, o: Self) -> Self {
unsafe { $func(self, o) }
}
}
};
}
impl_float!(from f64: __extenddftf2 __trunctfdf2);
impl_float!(from i64: __floatditf __fixtfdi);
impl_float!(from u64: __floatunditf __fixunstfdi);
impl_float!(from i32: __floatsitf __fixtfsi);
impl_float!(from u32: __floatunsitf __fixunstfsi);
impl_float!(op Add add __addtf3);
impl_float!(op Sub sub __subtf3);
impl_float!(op Mul mul __multf3);
impl_float!(op Div div __divtf3);
impl std::ops::Neg for f128 {
type Output = Self;
fn neg(self) -> Self::Output {
unsafe { __negtf3(self) }
}
}
#[link(name = "stdc++")]
#[allow(improper_ctypes)]
extern "C" {
fn __extenddftf2(a: f64) -> f128;
fn __trunctfdf2(a: f128) -> f64;
fn __addtf3(a: f128, b: f128) -> f128;
fn __subtf3(a: f128, b: f128) -> f128;
fn __multf3(a: f128, b: f128) -> f128;
fn __divtf3(a: f128, b: f128) -> f128;
fn __negtf3(a: f128) -> f128;
fn __fixtfsi(a: f128) -> i32;
fn __fixunstfsi(a: f128) -> u32;
fn __fixtfdi(a: f128) -> i64;
fn __fixunstfdi(a: f128) -> u64;
fn __floatsitf(a: i32) -> f128;
fn __floatunsitf(a: u32) -> f128;
fn __floatditf(a: i64) -> f128;
fn __floatunditf(a: u64) -> f128;
}GCD
fn gcd(mut a: i32, mut b: i32) -> i32 {
if a == 0 || b == 0 {
a + b
} else {
let az = a.trailing_zeros();
let bz = b.trailing_zeros();
let s = az.min(bz);
a >>= az;
b >>= bz;
while a != 0 {
let d = a - b;
(a, b) = (d.abs(), a.min(b));
a >>= d.trailing_zeros();
}
b << s
}
}
/// Returns 3-tuple `(g, x, y)` that satisfies ax + by = g = gcd(a, b).
fn egcd(mut a: i64, mut b: i64) -> (i64, i64, i64) {
let (mut x, mut y, mut x1, mut y1) = (1, 0, 0, 1);
while b != 0 {
let q = a / b;
(x, x1) = (x1, x - q * x1);
(y, y1) = (y1, y - q * y1);
(a, b) = (b, a - q * b);
}
(a, x, y)
}Linear Recurrence (Constant Coefficient)
Calculates nth element of linearly recurrent sequence {an}, given initial values and coefficients.
Time complexity: O(k³logn), where k is the length of coefficients.
const M: u64 = 998244353;
const M32: u32 = M as u32;
fn recurrence(init: impl AsRef<[u32]>, coef: impl AsRef<[u32]>, n: usize) -> u32 {
let init = init.as_ref();
let coef = coef.as_ref();
let k = init.len();
assert_eq!(k, coef.len());
if n < k {
return init[n];
}
let mut acc = vec![0; k * k];
for i in 0..k {
acc[i * k + i] = 1;
}
let mut mult = vec![0; k * (k - 1)];
for i in 1..k {
mult[(i - 1) * k + i] = 1;
}
mult.extend_from_slice(coef);
let mut aux = vec![0; k * k];
let t = (n - k + 1..=n).min_by_key(|t| t.count_ones()).unwrap();
for s in 0..64 - t.leading_zeros() {
if t & (1 << s) != 0 {
multiply(k, &acc, &mult, &mut aux);
(acc, aux) = (aux, acc);
}
multiply(k, &mult, &mult, &mut aux);
(mult, aux) = (aux, mult);
}
let mut result = 0;
for (&y, &x) in init.iter().zip(&acc[(n - t) * k..]) {
result += (y as u64 * x as u64 % M) as u32;
if result >= M32 {
result -= M32;
}
}
result
}
fn multiply(n: usize, a: &[u32], b: &[u32], c: &mut [u32]) {
for i in 0..n {
for j in 0..n {
c[i * n + j] = 0;
for k in 0..n {
c[i * n + j] += (a[i * n + k] as u64 * b[k * n + j] as u64 % M) as u32;
if c[i * n + j] >= M32 {
c[i * n + j] -= M32;
}
}
}
}
}Constant Division / Modulo
Speed up division and modulo calculation using precomputation.
reference: Daniel Lemire
#[derive(Copy, Clone)]
struct ModU32(u64, u32);
impl ModU32 {
const fn new(div: u32) -> Self {
Self((!0 / div as u64).wrapping_add(1), div)
}
const fn rem(&self, a: u32) -> u32 {
let low = self.0.wrapping_mul(a as u64);
((low as u128 * self.1 as u128) >> 64) as u32
}
const fn quot(&self, a: u32) -> u32 {
((self.0 as u128 * a as u128) >> 64) as u32
}
const fn quotrem(&self, a: u32) -> (u32, u32) {
let full = self.0 as u128 * a as u128;
let quot = full >> 64;
let rem = (full as u64 as u128 * self.1 as u128) >> 64;
(quot as u32, rem as u32)
}
const fn divisible(&self, a: u32) -> bool {
self.0.wrapping_mul(a as u64) <= self.0.wrapping_sub(1)
}
}
/// d != 0 && d != i32::MIN
#[derive(Copy, Clone)]
struct ModI32(u64, i32);
impl ModI32 {
const fn new(div: i32) -> Self {
let abs = div.abs();
let mut m = (!0 / abs as u64).wrapping_add(1);
if (abs & (abs - 1)) == 0 {
m += 1;
}
Self(m, div)
}
const fn rem(&self, a: i32) -> i32 {
let low = self.0.wrapping_mul(a as u64);
let high = ((low as u128 * self.1.abs() as u128) >> 64) as i32;
high - ((self.1 - 1) & (a >> 31))
}
const fn quot(&self, a: i32) -> i32 {
let mut high = ((self.0 as i128 * a as i128) >> 64) as u64;
if a < 0 {
high = high.wrapping_add(1);
}
if self.1 < 0 {
-(high as i32)
} else {
high as i32
}
}
const fn quotrem(&self, a: i32) -> (i32, i32) {
let full = self.0 as i128 * a as i128;
let mut high = (full >> 64) as u64;
if a < 0 {
high = high.wrapping_add(1);
}
let quot = if self.1 < 0 {
-(high as i32)
} else {
high as i32
};
let mut rem = ((full as u64 as u128 * self.1.abs() as u128) >> 64) as i32;
rem -= ((self.1 - 1) & a >> 31);
(quot, rem)
}
}
#[derive(Copy, Clone)]
struct ModU64(u128, u64);
impl ModU64 {
const fn new(div: u64) -> Self {
Self((!0u128 / div as u128).wrapping_add(1), div)
}
const fn multop(a: u128, b: u64) -> u64 {
let mut bottom = (a as u64 as u128) * b as u128;
bottom >>= 64;
let top = (a >> 64) * b as u128;
((bottom + top) >> 64) as u64
}
const fn rem(&self, a: u64) -> u64 {
let low = self.0.wrapping_mul(a as u128);
Self::multop(low, self.1)
}
const fn quot(&self, a: u64) -> u64 {
Self::multop(self.0, a)
}
const fn divisible(&self, a: u64) -> bool {
self.0.wrapping_mul(a as u128) <= self.0.wrapping_sub(1)
}
}
#[derive(Copy, Clone)]
struct BarrettU32(u64, u32);
impl BarrettU32 {
fn new(div: u32) -> Self {
Self(!0u64 / div as u64, div)
}
fn quot(&self, a: u64) -> u64 {
((self.0 as u128 * a as u128) >> 64) as u64
}
fn rem(&self, a: u64) -> u32 {
(a - self.quot(a) * self.1 as u64) as u32
}
}Number Theoretic Transform
struct NttParam {
m: ModU64,
u: u64,
ui: u64,
s: u32,
}
static NTT: &[NttParam] = &[NttParam {
m: ModU64::new(998244353),
u: 15311432,
ui: 469870224,
s: 23,
}, NttParam {
m: ModU64::new(1107296257),
u: 1087287097,
ui: 623044540,
s: 25,
}, NttParam {
m: ModU64::new(1711276033),
u: 969788637,
ui: 1790856,
s: 25,
}];
impl NttParam {
fn run(&self, a: &mut [u32], inv: bool) {
let n = a.len();
let s = n.leading_zeros() + 1;
for i in 0..n {
let r = i.reverse_bits() >> s;
if i < r {
a.swap(i, r);
}
}
let u = if inv { self.ui } else { self.u };
for k in 1..=n.trailing_zeros() {
let mut wlen = u;
for _ in k..self.s {
wlen = self.m.rem(wlen * wlen);
}
let kh = 1 << (k - 1);
for i in 0..(n + (1 << k) - 1) >> k {
let i = i << k;
let mut w = 1;
for j in 0..kh {
let u = a[i + j] as u64;
let v = self.m.rem(a[i + j + kh] as u64 * w);
let mut s = u + v;
if s >= self.m.1 {
s -= self.m.1;
}
a[i + j] = s as u32;
let mut d = u + self.m.1 - v;
if d >= self.m.1 {
d -= self.m.1;
}
a[i + j + kh] = d as u32;
w = self.m.rem(w * wlen);
}
}
}
if inv {
let p = self.m.1 as i64;
let ni = ((egcd(n as i64, p).1 % p + p) % p) as u64;
for x in a {
*x = self.m.rem(*x as u64 * ni) as u32;
}
}
}
}
fn egcd(mut a: i64, mut b: i64) -> (i64, i64, i64) {
let (mut x, mut y, mut x1, mut y1) = (1, 0, 0, 1);
while b != 0 {
let q = a / b;
(x, x1) = (x1, x - q * x1);
(y, y1) = (y1, y - q * y1);
(a, b) = (b, a - q * b);
}
(a, x, y)
}
#[derive(Copy, Clone)]
struct ModU64(u128, u64);
impl ModU64 {
const fn new(div: u64) -> Self {
Self((!0u128 / div as u128).wrapping_add(1), div)
}
fn multop(a: u128, b: u64) -> u64 {
let mut bottom = (a as u64 as u128) * b as u128;
bottom >>= 64;
let top = (a >> 64) * b as u128;
((bottom + top) >> 64) as u64
}
fn rem(&self, a: u64) -> u64 {
let low = self.0.wrapping_mul(a as u128);
Self::multop(low, self.1)
}
}Mulmod
Computes 64bit a times b mod c, returning quotient and remainder
fn mulmod(a: u64, b: u64, c: u64) -> (u64, u64) {
let (quot, rem);
unsafe {
std::arch::asm!(
"mul {b}",
"div {c}",
inout("rax") a => quot,
b = in(reg) b,
c = in(reg) c,
out("rdx") rem,
options(nomem, pure)
)
};
(quot, rem)
}
Convex Hull (Monotone Chain)
Returns points in convex hull in counterclockwise order.
Removes straight points. change is_le() to is_lt to contain straight points.
fn d((ax, ay): (i32, i32), (bx, by): (i32, i32)) -> f64 {
((bx - ax) as f64).hypot((by - ay) as f64)
}
fn sub((ax, ay): (i32, i32), (bx, by): (i32, i32)) -> (i32, i32) {
(ax - bx, ay - by)
}
fn cross((ax, ay): (i32, i32), (bx, by): (i32, i32)) -> i64 {
ax as i64 * by as i64 - ay as i64 * bx as i64
}
fn convex_hull(sorted_pts: &[(i32, i32)]) -> Vec<(i32, i32)> {
let mut hull = vec![];
let mut len = 0;
for &p in sorted_pts {
while len >= 2 && turn(hull[len - 2], hull[len - 1], p).is_le() {
hull.pop();
len -= 1;
}
hull.push(p);
len += 1;
}
hull.pop();
len -= 1;
let base = len + 2;
for &p in sorted_pts.iter().rev() {
while len >= base && turn(hull[len - 2], hull[len - 1], p).is_le() {
hull.pop();
len -= 1;
}
hull.push(p);
len += 1;
}
hull.pop();
hull
}Antipodal Points
An iterator that yields every antipodal pairs.
struct Antipodes<T> {
i: usize,
p: T,
q: T,
p1: T,
q1: T,
}
impl<'a, T> Iterator for Antipodes<std::iter::Peekable<T>>
where
T: Iterator<Item = &'a (i32, i32)> + Clone,
{
type Item = ((i32, i32), (i32, i32));
fn next(&mut self) -> Option<Self::Item> {
if self.i == 0 {
return None;
}
self.i -= 1;
let p = **self.p.peek()?;
let p1 = **self.p1.peek()?;
let q = **self.q.peek()?;
let q1 = **self.q1.peek()?;
if cross(p, p1, q) < cross(p, p1, q1) {
self.q = self.q1.clone();
Some((**self.p.peek()?, *self.q1.next()?))
} else {
self.p = self.p1.clone();
Some((*self.p1.next()?, **self.q.peek()?))
}
}
}
impl<'a> Antipodes<std::iter::Peekable<std::iter::Cycle<std::slice::Iter<'a, (i32, i32)>>>> {
fn from_hull(hull: &'a [(i32, i32)]) -> Self {
let mut base = hull.iter().cycle().peekable();
let mut p = base.clone();
base.next();
let mut p1 = base.clone();
let mut q = base.clone();
base.next();
let mut q1 = base;
let pp = **p.peek().unwrap();
let pp1 = **p1.peek().unwrap();
while cross(pp, pp1, **q.peek().unwrap()) < cross(pp, pp1, **q1.peek().unwrap()) {
q = q1.clone();
q1.next();
}
let i = hull.len();
Self { i, p, p1, q, q1 }
}
}
// Less => cw
// Greater => ccw
fn turn(a: (i32, i32), b: (i32, i32), c: (i32, i32)) -> std::cmp::Ordering {
cross(a, b, c).cmp(&0)
}
fn cross((ax, ay): (i32, i32), (bx, by): (i32, i32), (cx, cy): (i32, i32)) -> i64 {
let abx = (bx - ax) as i64;
let aby = (by - ay) as i64;
let bcx = (cx - bx) as i64;
let bcy = (cy - by) as i64;
abx * bcy - aby * bcx
}Segment Tree
struct SegTree<C, A, T> {
e: T,
v: Vec<T>,
offset: usize,
n: usize,
combine: C,
apply: A,
}
impl<C, A, T> SegTree<C, A, T>
where
T: Copy,
C: Fn(T, T) -> T,
{
fn new<I>(n: usize, e: T, combine: C, apply: A, init: I) -> Self
where
I: IntoIterator<Item = T>,
{
let offset = n;
let mut v = vec![e; offset];
v.extend(init.into_iter().take(n));
v.resize(offset * 2, e);
for i in (1..offset).rev() {
v[i] = combine(v[i << 1], v[i << 1 | 1]);
}
Self {
e,
v,
offset,
n,
combine,
apply,
}
}
fn query<B: std::ops::RangeBounds<usize>>(&self, range: B) -> T {
use std::ops::Bound::*;
let mut l = self.offset
+ match range.start_bound() {
Included(&x) => x,
Excluded(&x) => x + 1,
Unbounded => 0,
};
let mut r = self.offset
+ match range.end_bound() {
Included(&x) => x + 1,
Excluded(&x) => x,
Unbounded => self.n,
};
let mut lsum = self.e;
let mut rsum = self.e;
while l < r {
if l & 1 != 0 {
lsum = (self.combine)(lsum, self.v[l]);
l += 1;
}
if r & 1 != 0 {
r -= 1;
rsum = (self.combine)(self.v[r], rsum);
}
l >>= 1;
r >>= 1;
}
(self.combine)(lsum, rsum)
}
fn update<U>(&mut self, mut i: usize, u: U)
where
U: Copy,
A: Fn(&mut T, U),
{
i += self.offset;
(self.apply)(&mut self.v[i], u);
while i > 1 {
i >>= 1;
self.v[i] = (self.combine)(self.v[i << 1], self.v[i << 1 | 1]);
}
}
fn partition_point<P: Fn(T) -> bool>(&self, left: usize, pred: P) -> usize {
let mut p = left + self.offset;
let mut value = self.e;
loop {
if p & 1 != 0 {
if pred((self.combine)(value, self.v[p])) {
value = (self.combine)(value, self.v[p]);
p += 1;
} else {
break;
}
}
p >>= 1;
}
while p < self.offset {
p <<= 1;
if pred((self.combine)(value, self.v[p])) {
value = (self.combine)(value, self.v[p]);
p |= 1;
}
}
p - self.offset
}
}Splay Tree
Top-down Splay Tree, supports lazy operations.
#![allow(unused)] fn main() { use std::cmp::Ordering::*; trait SplayOp { type T; fn pull(_s: &mut Splay<Self>) {} fn push(_s: &mut Splay<Self>) {} } struct Splay<Op: SplayOp + ?Sized> { c: [Option<Box<Self>>; 2], k: u32, v: Op::T, } impl<Op: SplayOp> Splay<Op> { fn new(v: Op::T) -> Box<Self> { Box::new(Self { c: [None, None], k: 1, v, }) } fn pull(&mut self) { let l = self.c[0].as_ref().map_or(0, |c| c.k); let r = self.c[1].as_ref().map_or(0, |c| c.k); self.k = l + r + 1; Op::pull(self); } fn push(&mut self) { Op::push(self); } fn replace(&mut self, side: usize, tree: Option<Box<Self>>) -> Option<Box<Self>> { self.push(); let c = std::mem::replace(&mut self.c[side], tree); self.pull(); c } fn index(self: Box<Self>, k: usize) -> Box<Self> { let mut k = k as u32; self.find(|s| { let l = s.c[0].as_ref().map_or(0, |c| c.k); match k.cmp(&l) { Equal => 2, Less => 0, Greater => { k -= l + 1; 1 } } }) } fn find<F>(self: Box<Self>, mut f: F) -> Box<Self> where F: FnMut(&Self) -> usize, { let mut current = self; let mut tree = [vec![], vec![]]; loop { current.push(); let i = f(¤t); if let Some(mut c) = current.c.get_mut(i).and_then(Option::take) { c.push(); let j = f(&c); if let Some(cc) = c.c.get_mut(j).and_then(Option::take) { if i == j { let rotate = c.c[j ^ 1].take(); current.c[i] = rotate; current.pull(); c.c[j ^ 1] = Some(current); tree[j ^ 1].push(c); } else { tree[j].push(current); tree[i].push(c); } current = cc; } else { tree[i ^ 1].push(current); current = c; break; } } else { break; } } current.push(); let [mut left, mut right] = current.c; let [ltree, rtree] = tree; for mut l in ltree.into_iter().rev() { l.c[1] = left; l.pull(); left = Some(l); } for mut r in rtree.into_iter().rev() { r.c[0] = right; r.pull(); right = Some(r); } current.c = [left, right]; current.pull(); current } } }
Weighted Disjoint Set Union
struct Wdsu<W, C> {
up: Vec<(u32, W)>,
e: W,
combine: C,
}
impl<W, C> Wdsu<W, C>
where
W: Copy,
C: Fn(W, W) -> W,
{
fn new(n: usize, e: W, combine: C) -> Self {
Self {
up: (0..n as u32).map(|i| (i, e)).collect(),
e,
combine,
}
}
fn root(&mut self, u: usize) -> (usize, W) {
let mut r = u;
let mut w = self.e;
while r != self.up[r].0 as usize {
let (p, rw) = self.up[r];
let (g, pw) = self.up[p as usize];
self.up[r] = (g, (self.combine)(rw, pw));
w = (self.combine)(w, self.up[r].1);
r = g as usize;
}
(r, w)
}
fn attach(&mut self, p: usize, s: usize, w: W) {
self.up[s] = (p as u32, w);
}
}Graph
Graphs are represented in Compressed Sparse Row format. All edges on the graph are stored in a single array such that edges with the same starting vertex are connected like a linked list.
linkstores (previous edge index, endpoint) tuples.headstores indices of edge list heads (oru32::MAXif none).
#![allow(unused)] fn main() { struct Graph { head: Vec<u32>, link: Vec<(u32, u32)>, } impl Graph { fn new(v: usize, e: usize) -> Self { Self { head: vec![u32::MAX; v], link: Vec::with_capacity(e), } } fn connect(&mut self, from: usize, to: usize) { let prev = self.head[from]; self.head[from] = self.link.len() as u32; self.link.push((prev, to as u32)); } fn neighbor(&self, u: usize) -> Neighbor { Neighbor(self, self.head[u]) } } struct Neighbor<'a>(&'a Graph, u32); impl<'a> Iterator for Neighbor<'a> { // (edge index, endpoint) type Item = (usize, usize); fn next(&mut self) -> Option<Self::Item> { let e = self.1; let (next, endpoint) = *self.0.link.get(self.1 as usize)?; self.1 = next; Some((e as usize, endpoint as usize)) } } }
Heavy-Light Decomposition
Path traversing on heavy-light decomposition can be elegantly represented with Rust Iterator.
#![allow(unused)] fn main() { struct Hld { p: Vec<u32>, d: Vec<u32>, r: Vec<u32>, } impl Graph { fn hld(&self, root: usize) -> Hld { let n = self.head.len(); let mut d = vec![0; n]; let mut p = vec![0; n]; p[root] = root as u32; let mut bfs = vec![(root as u32, 0, 1, 0)]; let mut front = 0; let mut depth = 0; while front < bfs.len() { let back = bfs.len(); for i in front..back { let u = bfs[i].0 as usize; let pu = p[u]; d[u] = depth; for (_, v) in self.neighbor(u) { if v == pu as usize { continue; } bfs.push((v as u32, i as u32, 1, 0)); p[v] = u as u32; } } front = back; depth += 1; } let bfs = &mut bfs[..]; for i in (1..bfs.len()).rev() { let (_, p, sz, _) = bfs[i]; let (_, _, usz, umax) = &mut bfs[p as usize]; *usz += sz; *umax = sz.max(*umax); } let mut r = vec![0; n]; for i in 1..bfs.len() { let (v, p, sz, _) = bfs[i]; let (u, _, _, umax) = &mut bfs[p as usize]; if sz == *umax { r[v as usize] = *u; *umax = u32::MAX; bfs[i].0 = *u; } else { r[v as usize] = v; } } Hld { p, d, r } } } impl Hld { fn path(&self, u: usize, v: usize) -> HldPath { HldPath { hld: self, u, v } } fn lca(&self, mut u: usize, mut v: usize) -> usize { while self.r[u] != self.r[v] { if self.d[self.r[u] as usize] > self.d[self.r[v] as usize] { std::mem::swap(&mut u, &mut v); } v = self.p[self.r[v] as usize] as usize; } if self.d[u] > self.d[v] { v } else { u } } } struct HldPath<'a> { hld: &'a Hld, u: usize, v: usize, } impl Iterator for HldPath<'_> { // (upper, lower) type Item = (usize, usize); fn next(&mut self) -> Option<Self::Item> { if self.u == usize::MAX { return None; } let ru = self.hld.r[self.u] as usize; let mut rv = self.hld.r[self.v] as usize; if ru == rv { let ret = if self.hld.d[self.u] > self.hld.d[self.v] { (self.v, self.u) } else { (self.u, self.v) }; self.u = usize::MAX; Some(ret) } else { if self.hld.d[ru] > self.hld.d[rv] { (rv, self.u, self.v) = (ru, self.v, self.u); } self.v = self.hld.p[rv] as usize; Some((rv, self.v)) } } } }
Push-Relabel Flow
type Capacity = i32;
struct Flow {
head: Vec<u32>,
link: Vec<(u32, u32)>,
f: Vec<Capacity>,
x: Vec<Capacity>,
next: Vec<i32>,
last: Vec<i32>,
gn: Vec<i32>,
gp: Vec<i32>,
ff: Vec<u32>,
h: Vec<i32>,
high: i32,
high_gap: i32,
work: usize,
}
impl Flow {
fn new(v: usize, e: usize) -> Self {
Self {
head: vec![u32::MAX; v],
link: Vec::with_capacity(e),
f: Vec::with_capacity(e),
x: Vec::with_capacity(e),
gn: vec![],
gp: vec![],
h: vec![],
high: 0,
high_gap: 0,
work: 0,
next: vec![],
last: vec![],
ff: vec![],
}
}
fn connect(&mut self, from: usize, to: usize, cap: Capacity, directed: bool) {
let p = self.head[from];
self.head[from] = self.link.len() as u32;
self.link.push((p, to as u32));
self.f.push(cap);
let p = self.head[to];
self.head[to] = self.link.len() as u32;
self.link.push((p, from as u32));
self.f.push(if directed { 0 } else { cap });
}
fn push_last(&mut self, h: i32, u: usize) {
self.next[u] = self.last[h as usize];
self.last[h as usize] = u as i32;
}
fn update_h(&mut self, u: usize, mut new_h: i32) {
if self.h[u] as usize != self.head.len() {
self.gn[self.gp[u] as usize] = self.gn[u];
self.gp[self.gn[u] as usize] = self.gp[u];
}
self.h[u] = new_h as i32;
if new_h == self.head.len() as i32 {
return;
}
self.high_gap = self.high_gap.max(new_h);
if self.x[u] > 0 {
self.high = self.high.max(new_h);
self.push_last(new_h, u);
}
new_h += self.head.len() as i32;
self.gn[u] = self.gn[new_h as usize];
self.gp[u] = new_h as i32;
self.gn[new_h as usize] = u as i32;
self.gp[self.gn[u] as usize] = u as i32;
}
fn global_relabel(&mut self, t: usize) {
self.work = 0;
self.h.clear();
self.h
.extend(std::iter::repeat(self.head.len() as i32).take(self.head.len()));
self.last.clear();
self.last.resize(self.head.len() + 1, -1);
self.gn[..self.head.len()]
.iter_mut()
.enumerate()
.for_each(|(i, v)| *v = i as i32);
self.gp[..self.head.len()]
.iter_mut()
.enumerate()
.for_each(|(i, v)| *v = i as i32);
self.h[t] = 0;
let mut q = vec![t as u32];
let mut i = 0;
while let Some(&u) = q.get(i) {
let mut e = self.head[u as usize];
while let Some(&(next, v)) = self.link.get(e as usize) {
if self.h[v as usize] == self.head.len() as i32 && self.f[(e ^ 1) as usize] > 0 {
q.push(v);
self.update_h(v as usize, self.h[u as usize] + 1);
}
e = next;
}
self.high_gap = self.h[u as usize];
self.high = self.high_gap;
i += 1;
}
}
fn push(&mut self, u: usize, v: usize, e: usize) {
let df = self.x[u].min(self.f[e]);
if df <= 0 {
return;
}
if self.x[v] == 0 {
self.push_last(self.h[v], v as usize);
}
self.f[e] -= df;
self.f[e ^ 1] += df;
self.x[u] -= df;
self.x[v] += df;
}
fn discharge(&mut self, u: usize) {
let mut new_h = self.head.len() as i32;
let mut e = self.ff[u];
while let Some(&(next, v)) = self.link.get(e as usize) {
if self.f[e as usize] > 0 {
if self.h[u] == self.h[v as usize] + 1 {
self.push(u, v as usize, e as usize);
if self.x[u] <= 0 {
self.ff[u] = e;
return;
}
} else {
new_h = new_h.min(self.h[v as usize] + 1);
}
}
e = next;
}
let mut e = self.head[u];
while e != self.ff[u] {
let (next, v) = self.link[e as usize];
if self.f[e as usize] > 0 {
if self.h[u] == self.h[v as usize] + 1 {
self.push(u, v as usize, e as usize);
if self.x[u] <= 0 {
self.ff[u] = e;
return;
}
} else {
new_h = new_h.min(self.h[v as usize] + 1);
}
}
e = next;
}
self.work += 1;
let h = self.h[u] + self.head.len() as i32;
if self.gn[self.gn[h as usize] as usize] != h {
self.update_h(u, new_h);
} else {
let old = self.h[u];
for h in old..=self.high_gap {
let mut u = self.gn[(h + self.head.len() as i32) as usize] as usize;
while u < self.head.len() {
self.h[u] = self.head.len() as i32;
u = self.gn[u] as usize;
}
let uh = h + self.head.len() as i32;
self.gp[uh as usize] = uh;
self.gn[uh as usize] = uh;
}
self.high_gap = old - 1;
}
}
fn run(&mut self, s: usize, t: usize) -> Capacity {
self.ff.clear();
self.ff.extend_from_slice(&self.head);
self.x.clear();
self.x.resize(self.head.len(), 0);
self.x[s] = Capacity::MAX;
self.x[t] = -Capacity::MAX;
self.gn.clear();
self.gn.resize(self.head.len() * 2, 0);
self.gp.clear();
self.gp.resize(self.head.len() * 2, 0);
self.next.clear();
self.next.resize(self.head.len(), 0);
self.global_relabel(t);
let mut e = self.head[s];
while let Some(&(next, v)) = self.link.get(e as usize) {
self.push(s, v as usize, e as usize);
e = next;
}
while self.high >= 0 {
while self.last[self.high as usize] != -1 {
let u = self.last[self.high as usize];
self.last[self.high as usize] = self.next[u as usize];
if self.h[u as usize] == self.high {
self.discharge(u as usize);
if self.work > 4 * self.head.len() {
self.global_relabel(t);
}
}
}
self.high -= 1;
}
self.x[t] + i32::MAX
}
}Min Heap
pop()- Returns (inserted index, key) pair.
struct MinHeap<K> {
k: Vec<(u32, K)>,
pos: Vec<u32>,
max: usize,
}
impl<K> MinHeap<K>
where
K: PartialOrd + Copy,
{
fn new() -> Self {
Self {
k: vec![],
pos: vec![],
max: 0,
}
}
fn len(&self) -> usize {
self.k.len()
}
fn push(&mut self, k: K) {
let i = self.max;
let len = self.len() as u32;
self.max += 1;
self.k.push((i as u32, k));
self.pos.push(len);
self.up(i);
}
fn pop(&mut self) -> Option<(usize, K)> {
if self.len() == 0 {
return None;
}
let (i, k) = self.k.swap_remove(0);
if self.len() > 0 {
self.pos[self.k[0].0 as usize] = 0;
self.down(0);
}
Some((i as usize, k))
}
fn replace(&mut self, i: usize, k: K) -> K {
use std::cmp::Ordering::*;
let pos = self.pos[i] as usize;
let ordering = self.k[pos].1.partial_cmp(&k).unwrap();
let prev = std::mem::replace(&mut self.k[pos].1, k);
match ordering {
Greater => self.up(pos),
Less => self.down(pos),
Equal => {}
}
prev
}
fn up(&mut self, mut pos: usize) {
while pos > 0 {
let parent = (pos - 1) / 2;
if self.k[pos].1 >= self.k[parent].1 {
break;
}
self.k.swap(pos, parent);
self.pos[self.k[pos].0 as usize] = pos as u32;
pos = parent;
}
self.pos[self.k[pos].0 as usize] = pos as u32;
}
fn down(&mut self, mut pos: usize) {
let len = self.len();
loop {
let child = pos * 2 + 1;
if child >= len {
break;
}
let less = if child + 1 == len || self.k[child].1 <= self.k[child + 1].1 {
child
} else {
child + 1
};
if self.k[pos].1 <= self.k[less].1 {
break;
}
self.k.swap(less, pos);
self.pos[self.k[pos].0 as usize] = pos as u32;
pos = less;
}
self.pos[self.k[pos].0 as usize] = pos as u32;
}
}
impl<K> FromIterator<K> for MinHeap<K>
where
K: PartialOrd + Copy,
{
fn from_iter<T: IntoIterator<Item = K>>(iter: T) -> Self {
let k: Vec<_> = iter
.into_iter()
.enumerate()
.map(|(i, v)| (i as u32, v))
.collect();
let len = k.len();
let pos = (0..len as u32).collect();
let mut inst = Self { k, pos, max: len };
for i in (0..inst.len()).rev() {
inst.down(i);
}
inst
}
}
Example
Single Source Shortest Path (Dijkstra's)
fn sssp(g: &Graph, w: &[u32], s: usize) -> Vec<u32> {
let mut d = vec![u32::MAX; g.head.len()];
d[s] = 0;
let mut heap = d.iter().map(|&d| d).collect::<MinHeap<_>>();
while let Some((u, du)) = heap.pop() {
for (e, v) in g.neighbor(u) {
let dv = du + w[e >> 1] as u32;
if d[v] > dv {
d[v] = dv;
heap.replace(v, dv);
}
}
}
d
}
Shallowest Decomposition Tree
fn decompose_tree(t: &Graph, root: usize) -> (usize, Graph) {
let n = t.head.len();
let mut result = Graph::new(n, n - 1);
let mut stack = Graph::new(64 - n.leading_zeros() as usize, n);
let mut extract_chain = |stack: &mut Graph, mut labels: u32, mut u: usize| {
while labels != 0 {
let label = labels.ilog2();
labels ^= 1 << label;
let back = stack.head[label as usize];
let (next, v) = stack.link[back as usize];
stack.head[label as usize] = next;
result.connect(u, v as usize);
u = v as usize;
}
};
let mut forbidden = vec![0; n];
let mut dfs = vec![(root as u32, u32::MAX, t.head[root], 0u32, 0u32)];
let mut last_f = 0;
while let Some((u, pu, eu, f1, f2)) = dfs.last_mut() {
if let Some(&(ev, v)) = t.link.get(*eu as usize) {
let now = *eu;
*eu = ev;
if now == *pu {
continue;
}
dfs.push((v, now ^ 1, t.head[v as usize], 0, 0));
} else {
let fu = (*f1 | ((1 << (*f2 * 2 + 1).ilog2()) - 1)) + 1;
forbidden[*u as usize] = fu;
last_f = fu;
let lu = fu.trailing_zeros();
stack.connect(lu as usize, *u as usize);
for (_, v) in t.neighbor(*u as usize) {
extract_chain(&mut stack, forbidden[v] & ((1 << lu) - 1), *u as usize);
}
dfs.pop();
let fv = fu;
if let Some((_, _, _, f1, f2)) = dfs.last_mut() {
*f2 |= *f1 & fv;
*f1 |= fv;
}
}
}
let maxl = last_f.ilog2() as usize;
let back = stack.head[maxl];
let (next, root) = stack.link[back as usize];
stack.head[maxl] = next;
extract_chain(&mut stack, last_f & ((1 << maxl) - 1), root as usize);
(root as usize, result)
}KMP
struct Kmp<'a, I, T> {
pi: Vec<u32>,
haystack: I,
needle: &'a [T],
i: usize,
}
impl<'a, I, T> Kmp<'a, I, T> {
fn new<H>(haystack: H, needle: &'a [T], pi: Vec<u32>) -> Self
where
H: IntoIterator<IntoIter = I>,
{
Self {
pi,
haystack: haystack.into_iter(),
needle,
i: 0,
}
}
}
impl<I, T, B> Iterator for Kmp<'_, I, T>
where
T: PartialEq,
B: std::borrow::Borrow<T>,
I: Iterator<Item = B>,
{
type Item = usize;
fn next(&mut self) -> Option<Self::Item> {
let c = self.haystack.next()?;
loop {
if self.needle.get(self.i) == Some(c.borrow()) {
self.i += 1;
break;
} else if let Some(&pi) = self.pi.get(self.i.wrapping_sub(1)) {
self.i = pi as usize;
} else {
break;
}
}
Some(self.i)
}
}Example
let haystack = b"AABABAABABAA";
let needle = b"ABA";
let mut kmp = Kmp::new(&needle[1..], needle, vec![0]);
while let Some(pi) = kmp.next() {
kmp.pi.push(pi as u32);
}
for (i, len) in Kmp::new(haystack, needle, kmp.pi).enumerate() {
if len > 0 {
println!("Partial match [{}, {}]", i + 1 - len, i);
}
}Regex
Use POSIX library function.
struct Regexp([u64; 8]);
extern "C" {
fn regcomp(preg: *mut u64, pattern: *const u8, flags: i32);
fn regexec(preg: *const u64, s: *const u8, n_expr: usize, m: *mut i32, flags: i32) -> i32;
}
impl Regexp {
fn new(pattern: &str, case_insensitive: bool) -> Self {
let flag = if case_insensitive { 3 } else { 1 };
let mut tmp = pattern.as_bytes().to_vec();
tmp.push(0);
let mut v = [0; 8];
unsafe { regcomp(v.as_mut_ptr(), tmp.as_ptr(), flag) };
Self(v)
}
fn find(&self, text: &str) -> Option<(usize, usize)> {
let mut text = text.as_bytes().to_vec();
text.push(0);
let mut m = [0; 2];
let result = unsafe { regexec(self.0.as_ptr(), text.as_ptr(), 1, m.as_mut_ptr(), 0) };
if result == 0 {
Some((m[0] as usize, m[1] as usize))
} else {
None
}
}
}