#656185 (Rust) No.1712 Read and Pile

提出ソース

結果

問題	No.1712 Read and Pile
ユーザー	akakimidori
提出日時	2021-05-11 20:42:51
言語	Rust (1.77.0 + proconio)
結果	AC
実行時間	445 ms / 2,000 ms
コード長	10,338 bytes
コンパイル時間	12,954 ms
コンパイル使用メモリ	401,712 KB
実行使用メモリ	22,384 KB
最終ジャッジ日時	2024-09-17 16:54:24
合計ジャッジ時間	24,099 ms
ジャッジサーバーID （参考情報）	judge5 / judge4

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 6,816 KB
testcase_01	AC	1 ms 6,816 KB
testcase_02	AC	1 ms 6,940 KB
testcase_03	AC	1 ms 6,940 KB
testcase_04	AC	1 ms 6,944 KB
testcase_05	AC	1 ms 6,940 KB
testcase_06	AC	1 ms 6,940 KB
testcase_07	AC	1 ms 6,940 KB
testcase_08	AC	261 ms 11,660 KB
testcase_09	AC	372 ms 20,360 KB
testcase_10	AC	242 ms 11,728 KB
testcase_11	AC	302 ms 20,484 KB
testcase_12	AC	208 ms 11,968 KB
testcase_13	AC	268 ms 11,928 KB
testcase_14	AC	237 ms 12,012 KB
testcase_15	AC	309 ms 11,596 KB
testcase_16	AC	313 ms 12,320 KB
testcase_17	AC	281 ms 12,012 KB
testcase_18	AC	289 ms 11,688 KB
testcase_19	AC	263 ms 11,644 KB
testcase_20	AC	265 ms 12,120 KB
testcase_21	AC	338 ms 20,508 KB
testcase_22	AC	280 ms 20,292 KB
testcase_23	AC	369 ms 22,164 KB
testcase_24	AC	333 ms 22,200 KB
testcase_25	AC	381 ms 22,384 KB
testcase_26	AC	445 ms 22,192 KB
testcase_27	AC	442 ms 22,168 KB
testcase_28	AC	214 ms 12,272 KB
testcase_29	AC	346 ms 21,880 KB
testcase_30	AC	266 ms 21,180 KB
testcase_31	AC	156 ms 22,116 KB
testcase_32	AC	110 ms 21,760 KB
testcase_33	AC	259 ms 21,148 KB
testcase_34	AC	254 ms 21,064 KB
testcase_35	AC	397 ms 21,544 KB
testcase_36	AC	444 ms 21,892 KB
testcase_37	AC	393 ms 21,400 KB
testcase_38	AC	1 ms 6,940 KB
testcase_39	AC	4 ms 6,940 KB
testcase_40	AC	5 ms 6,940 KB
testcase_41	AC	123 ms 10,624 KB
testcase_42	AC	112 ms 10,836 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

// ---------- begin Lazy Segment Tree ----------
pub trait TE {
    type T: Clone;
    type E: Clone;
    fn fold(l: &Self::T, r: &Self::T) -> Self::T;
    fn eval(x: &Self::T, f: &Self::E) -> Self::T;
    fn merge(g: &Self::E, h: &Self::E) -> Self::E;
    fn e() -> Self::T;
    fn id() -> Self::E;
}

pub struct LazySegmentTree<R: TE> {
    size: usize,
    bit: usize,
    a: Vec<(R::T, R::E)>,
}

impl <R: TE> LazySegmentTree<R> {
    pub fn new(n: usize) -> LazySegmentTree<R> {
        let size = n.next_power_of_two();
        let bit = size.trailing_zeros() as usize;
        LazySegmentTree {
            size: size,
            bit: bit,
            a: vec![(R::e(), R::id()); 2 * size],
        }
    }
    pub fn build_by(z: &[R::T]) -> LazySegmentTree<R> {
        let mut seg = LazySegmentTree::<R>::new(z.len());
        for (a, z) in seg.a[seg.size..].iter_mut().zip(z.iter()) {
            a.0 = z.clone();
        }
        let a = &mut seg.a;
        for i in (1..seg.size).rev() {
            a[i].0 = R::fold(&a[2 * i].0, &a[2 * i + 1].0);
        }
        seg
    }
    fn apply(&mut self, x: usize, op: &R::E) {
        let node = &mut self.a[x];
        node.0 = R::eval(&node.0, op);
        node.1 = R::merge(&node.1, op);
    }
    fn propagate(&mut self, x: usize) {
        let mut op = R::id();
        std::mem::swap(&mut op, &mut self.a[x].1);
        self.apply(2 * x, &op);
        self.apply(2 * x + 1, &op);
    }
    fn propagate_range(&mut self, l: usize, r: usize) {
        let x = l + self.size;
        let y = r + self.size;
        let mut k = self.bit;
        while (x >> k) == (y >> k) {
            self.propagate(x >> k);
            k -= 1;
        }
        for i in ((x.trailing_zeros() as usize + 1)..=k).rev() {
            self.propagate(x >> i);
        }
        for i in ((y.trailing_zeros() as usize + 1)..=k).rev() {
            self.propagate(y >> i);
        }
    }
    fn save_range(&mut self, l: usize, r: usize) {
        let mut x = l + self.size;
        let mut y = r + self.size;
        let mut p = (x & 1) == 1;
        let mut q = (y & 1) == 1;
        x >>= 1;
        y >>= 1;
        while 0 < x && x < y {
            if p {
                self.a[x].0 = R::fold(&self.a[2 * x].0, &self.a[2 * x + 1].0);
            }
            if q {
                self.a[y].0 = R::fold(&self.a[2 * y].0, &self.a[2 * y + 1].0);
            }
            p |= (x & 1) == 1;
            q |= (y & 1) == 1;
            x >>= 1;
            y >>= 1;
        }
        while 0 < x {
            self.a[x].0 = R::fold(&self.a[2 * x].0, &self.a[2 * x + 1].0);
            x >>= 1;
        }
    }
    pub fn set_at(&mut self, x: usize, v: R::T) {
        assert!(x < self.size);
        let mut pos = x + self.size;
        for i in (1..=self.bit).rev() {
            self.propagate(pos >> i);
        }
        self.a[pos].0 = v;
        pos >>= 1;
        while pos > 0 {
            self.a[pos].0 = R::fold(&self.a[2 * pos].0, &self.a[2 * pos + 1].0);
            pos >>= 1;
        }
    }
    pub fn update(&mut self, l: usize, r: usize, op: R::E) {
        self.propagate_range(l, r);
        let mut x = l + self.size;
        let mut y = r + self.size;
        while x < y {
            if x & 1 == 1 {
                self.apply(x, &op);
                x += 1;
            }
            if y & 1 == 1 {
                y -= 1;
                self.apply(y, &op);
            }
            x >>= 1;
            y >>= 1;
        }
        self.save_range(l, r);
    }
    pub fn find(&mut self, l: usize, r: usize) -> R::T {
        self.propagate_range(l, r);
        let mut x = l + self.size;
        let mut y = r + self.size;
        let mut p = R::e();
        let mut q = R::e();
        while x < y {
            if x & 1 == 1 {
                p = R::fold(&p, &self.a[x].0);
                x += 1;
            }
            if y & 1 == 1 {
                y -= 1;
                q = R::fold(&self.a[y].0, &q);
            }
            x >>= 1;
            y >>= 1;
        }
        R::fold(&p, &q)
    }
}
// ---------- end Lazy Segment Tree ----------
// ---------- begin ModInt ----------
mod modint {

    #[allow(dead_code)]
    pub struct Mod;
    impl ConstantModulo for Mod {
        const MOD: u32 = 998_244_353;
    }

    use std::marker::*;
    use std::ops::*;

    pub trait Modulo {
        fn modulo() -> u32;
    }

    pub trait ConstantModulo {
        const MOD: u32;
    }

    impl<T> Modulo for T
    where
        T: ConstantModulo,
    {
        fn modulo() -> u32 {
            T::MOD
        }
    }

    pub struct ModInt<T>(pub u32, PhantomData<T>);

    impl<T> Clone for ModInt<T> {
        fn clone(&self) -> Self {
            ModInt::new_unchecked(self.0)
        }
    }

    impl<T> Copy for ModInt<T> {}

    impl<T: Modulo> Add for ModInt<T> {
        type Output = ModInt<T>;
        fn add(self, rhs: Self) -> Self::Output {
            let mut d = self.0 + rhs.0;
            if d >= T::modulo() {
                d -= T::modulo();
            }
            ModInt::new_unchecked(d)
        }
    }

    impl<T: Modulo> AddAssign for ModInt<T> {
        fn add_assign(&mut self, rhs: Self) {
            *self = *self + rhs;
        }
    }

    impl<T: Modulo> Sub for ModInt<T> {
        type Output = ModInt<T>;
        fn sub(self, rhs: Self) -> Self::Output {
            let mut d = T::modulo() + self.0 - rhs.0;
            if d >= T::modulo() {
                d -= T::modulo();
            }
            ModInt::new_unchecked(d)
        }
    }

    impl<T: Modulo> SubAssign for ModInt<T> {
        fn sub_assign(&mut self, rhs: Self) {
            *self = *self - rhs;
        }
    }

    impl<T: Modulo> Mul for ModInt<T> {
        type Output = ModInt<T>;
        fn mul(self, rhs: Self) -> Self::Output {
            let v = self.0 as u64 * rhs.0 as u64 % T::modulo() as u64;
            ModInt::new_unchecked(v as u32)
        }
    }

    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 = ModInt<T>;
        fn neg(self) -> Self::Output {
            if self.0 == 0 {
                Self::zero()
            } else {
                Self::new_unchecked(T::modulo() - self.0)
            }
        }
    }

    impl<T> std::fmt::Display for ModInt<T> {
        fn fmt<'a>(&self, f: &mut std::fmt::Formatter<'a>) -> std::fmt::Result {
            write!(f, "{}", self.0)
        }
    }

    impl<T: Modulo> From<usize> for ModInt<T> {
        fn from(val: usize) -> ModInt<T> {
            ModInt::new_unchecked((val % T::modulo() as usize) as u32)
        }
    }

    impl<T: Modulo> From<u64> for ModInt<T> {
        fn from(val: u64) -> ModInt<T> {
            ModInt::new_unchecked((val % T::modulo() as u64) as u32)
        }
    }

    #[allow(dead_code)]
    impl<T> ModInt<T> {
        pub fn new_unchecked(d: u32) -> Self {
            ModInt(d, PhantomData)
        }
        pub fn zero() -> Self {
            ModInt::new_unchecked(0)
        }
        pub fn one() -> Self {
            ModInt::new_unchecked(1)
        }
        pub fn is_zero(&self) -> bool {
            self.0 == 0
        }
    }

    #[allow(dead_code)]
    impl<T: Modulo> ModInt<T> {
        pub fn new(d: u32) -> Self {
            ModInt::new_unchecked(d % T::modulo())
        }
        pub fn pow(&self, mut n: u64) -> Self {
            let mut t = Self::one();
            let mut s = *self;
            while n > 0 {
                if n & 1 == 1 {
                    t *= s;
                }
                s *= s;
                n >>= 1;
            }
            t
        }
        pub fn inv(&self) -> Self {
            assert!(self.0 != 0);
            self.pow(T::modulo() as u64 - 2)
        }
    }
}
// ---------- end ModInt ----------

use modint::*;
type M = ModInt<Mod>;

const MIN_N: usize = 1;
const MAX_N: usize = 200_000;
const MIN_M: usize = 1;
const MAX_M: usize = 200_000;

fn read() -> (usize, Vec<i32>) {
    let mut s = String::new();
    use std::io::Read;
    std::io::stdin().read_to_string(&mut s).unwrap();
    let mut it = s.trim().split_whitespace();
    let mut next = |l: i32, r: i32| -> i32 {
        let v = it.next().unwrap().parse::<i32>().unwrap();
        assert!(l <= v && v <= r);
        v
    };
    let n = next(MIN_N as i32, MAX_N as i32) as usize;
    let m = next(MIN_M as i32, MAX_M as i32) as usize;
    let mut a = vec![0; m];
    for a in a.iter_mut() {
        *a = next(-1, n as i32);
        assert!(*a != 0);
    }
    assert!(it.next().is_none());
    (n, a)
}

struct R;
impl TE for R {
    type T = (M, M);
    type E = (M, M);
    fn fold(l: &Self::T, r: &Self::T) -> Self::T {
        (l.0 + r.0, l.1 + r.1)
    }
    fn eval(x: &Self::T, f: &Self::E) -> Self::T {
        (x.0 * f.0 + f.1 * x.1, x.1)
    }
    fn merge(g: &Self::E, h: &Self::E) -> Self::E {
        (g.0 * h.0, h.0 * g.1 + h.1)
    }
    fn e() -> Self::T {
        (M::zero(), M::zero())
    }
    fn id() -> Self::E {
        (M::one(), M::zero())
    }
}

fn solve(n: usize, a: &[i32]) -> M {
    if n == 1 {
        return M::from(a.len());
    }
    let m = a.len();
    let mut seg = LazySegmentTree::<R>::new(n + m + 1);
    let mut memo = vec![(m, 0); n];
    for i in 0..n {
        memo[i].0 += i;
        seg.set_at(memo[i].0, (M::one(), M::one()));
    }
    let zero = M::zero();
    let one = M::one();
    let p = M::from(n).inv();
    let q = M::from(n - 2) * p;
    let add = M::from(n - 1) * M::new(2).inv();
    let mut po = m - 1;
    let mut cnt = 0;
    let mut ans = M::from(m);
    for &a in a.iter() {
        if a == -1 {
            ans += add;
            seg.update(0, n + m, (q, p));
            cnt += 1;
        } else {
            let memo = &mut memo[a as usize - 1];
            let (s, c) = seg.find(po, memo.0);
            ans += s;
            ans += (M::from(n - 1) - c) * p * (one - q.pow(cnt - memo.1)) * (one - q).inv();
            seg.set_at(memo.0, (zero, zero));
            *memo = (po, cnt);
            po -= 1;
            seg.set_at(memo.0, (one, one));
        }
    }
    ans
}

fn main() {
    let (n, a) = read();
    let x = solve(n, &a);
    println!("{}", x);
}

yukicoder

結果

テストケース

ソースコード