【Rust】ACL の遅延セグ木をベースに双対セグ木を作ってみた

ACL（AtCoder Library）には通常のセグ木と遅延セグ木はありますが、双対セグ木はありません（2024年5月時点）。

遅延セグ木は区間作用・区間取得ができるデータ構造なのに対して、双対セグ木は区間作用・1点取得ができるデータ構造です。 双対セグ木は遅延セグ木と比べると区間取得ができないです。一方、双対セグ木には以下のメリットがあります。

• 乗せるモノイドの定義が簡潔になる。
  • 遅延セグ木では作用とデータの両方に対してモノイドを定義する必要があったのに対して、双対セグ木は作用に対してのみモノイドを定義すればよいためです。
• 遅延セグ木と比べて双対セグ木の方が定数倍で速い

つまり、1点更新で十分であれば双対セグ木は有用です。 そこで、ACL の遅延セグ木をベースに双対セグ木を作成しました。 基本的には、遅延セグ木から機能を削る形で作成しました。

「ACL の遅延セグ木のコード」と「作成した双対セグ木のコード」との diff

（作成した双対セグ木のコードは本記事の後半にあります）

diff を見ると、双対セグ木のコードの作り方が大体わかります。diff を要約すると以下の通りです。

• 作用を受けるデータに対してはモノイドが不要になった（作用に対してのみモノイドを定義すればよくなった）

• データはセグ木の葉に対してのみ持つようにした（遅延セグ木では、セグ木の各ノードに対してデータを載せていた）

  • データを持つ配列の長さは、遅延セグ木では 2 * size だったのに対して、双対セグ木では size になっている
• 区間取得の関数(prod など)を削除した
• (など)

実際のコードの diff は以下の通りです。

--- acl_lazy_segtree.rs  2024-05-05 22:01:52.807678542 +0900
+++ dual_segtree.rs   2024-05-05 22:27:35.097668299 +0900
@@ -1,148 +1,78 @@
-use crate::internal_bit::ceil_pow2;
-use crate::Monoid;
+use std::ops::{Bound, RangeBounds};
+
+fn ceil_pow2(n: u32) -> u32 {
+    32 - n.saturating_sub(1).leading_zeros()
+}
 
 pub trait MapMonoid {
-    type M: Monoid;
     type F: Clone;
-    // type S = <Self::M as Monoid>::S;
-    fn identity_element() -> <Self::M as Monoid>::S {
-        Self::M::identity()
-    }
-    fn binary_operation(
-        a: &<Self::M as Monoid>::S,
-        b: &<Self::M as Monoid>::S,
-    ) -> <Self::M as Monoid>::S {
-        Self::M::binary_operation(a, b)
-    }
+    type S: Clone;
     fn identity_map() -> Self::F;
-    fn mapping(f: &Self::F, x: &<Self::M as Monoid>::S) -> <Self::M as Monoid>::S;
+    fn mapping(f: &Self::F, x: &Self::S) -> Self::S;
     fn composition(f: &Self::F, g: &Self::F) -> Self::F;
 }
 
-impl<F: MapMonoid> Default for LazySegtree<F> {
+impl<F: MapMonoid> Default for DualSegtree<F>
+where
+    F::S: Default,
+{
     fn default() -> Self {
         Self::new(0)
     }
 }
-impl<F: MapMonoid> LazySegtree<F> {
-    pub fn new(n: usize) -> Self {
-        vec![F::identity_element(); n].into()
+impl<F: MapMonoid> DualSegtree<F> {
+    pub fn new(n: usize) -> Self
+    where
+        F::S: Default,
+    {
+        vec![F::S::default(); n].into()
     }
 }
-impl<F: MapMonoid> From<Vec<<F::M as Monoid>::S>> for LazySegtree<F> {
-    fn from(v: Vec<<F::M as Monoid>::S>) -> Self {
+
+impl<F: MapMonoid> From<Vec<F::S>> for DualSegtree<F>
+where
+    F::S: Default,
+{
+    fn from(v: Vec<F::S>) -> Self {
         let n = v.len();
         let log = ceil_pow2(n as u32) as usize;
         let size = 1 << log;
-        let mut d = vec![F::identity_element(); 2 * size];
+        let mut d = vec![F::S::default(); size];
         let lz = vec![F::identity_map(); size];
-        d[size..(size + n)].clone_from_slice(&v);
-        let mut ret = LazySegtree {
+        d[..n].clone_from_slice(&v);
+        DualSegtree {
             n,
             size,
             log,
             d,
             lz,
-        };
-        for i in (1..size).rev() {
-            ret.update(i);
         }
-        ret
     }
 }
 
-impl<F: MapMonoid> LazySegtree<F> {
-    pub fn set(&mut self, mut p: usize, x: <F::M as Monoid>::S) {
+impl<F: MapMonoid> DualSegtree<F> {
+    pub fn set(&mut self, p: usize, x: F::S) {
         assert!(p < self.n);
-        p += self.size;
         for i in (1..=self.log).rev() {
-            self.push(p >> i);
+            self.push((p + self.size) >> i);
         }
         self.d[p] = x;
-        for i in 1..=self.log {
-            self.update(p >> i);
-        }
     }
 
-    pub fn get(&mut self, mut p: usize) -> <F::M as Monoid>::S {
+    pub fn get(&mut self, p: usize) -> F::S {
         assert!(p < self.n);
-        p += self.size;
         for i in (1..=self.log).rev() {
-            self.push(p >> i);
+            self.push((p + self.size) >> i);
         }
         self.d[p].clone()
     }
 
-    pub fn prod<R>(&mut self, range: R) -> <F::M as Monoid>::S
-    where
-        R: RangeBounds<usize>,
-    {
-        // Trivial optimization
-        if range.start_bound() == Bound::Unbounded && range.end_bound() == Bound::Unbounded {
-            return self.all_prod();
-        }
-
-        let mut r = match range.end_bound() {
-            Bound::Included(r) => r + 1,
-            Bound::Excluded(r) => *r,
-            Bound::Unbounded => self.n,
-        };
-        let mut l = match range.start_bound() {
-            Bound::Included(l) => *l,
-            Bound::Excluded(l) => l + 1,
-            // TODO: There are another way of optimizing [0..r)
-            Bound::Unbounded => 0,
-        };
-
-        assert!(l <= r && r <= self.n);
-        if l == r {
-            return F::identity_element();
-        }
-
-        l += self.size;
-        r += self.size;
-
-        for i in (1..=self.log).rev() {
-            if ((l >> i) << i) != l {
-                self.push(l >> i);
-            }
-            if ((r >> i) << i) != r {
-                self.push(r >> i);
-            }
-        }
-
-        let mut sml = F::identity_element();
-        let mut smr = F::identity_element();
-        while l < r {
-            if l & 1 != 0 {
-                sml = F::binary_operation(&sml, &self.d[l]);
-                l += 1;
-            }
-            if r & 1 != 0 {
-                r -= 1;
-                smr = F::binary_operation(&self.d[r], &smr);
-            }
-            l >>= 1;
-            r >>= 1;
-        }
-
-        F::binary_operation(&sml, &smr)
-    }
-
-    pub fn all_prod(&self) -> <F::M as Monoid>::S {
-        self.d[1].clone()
-    }
-
-    pub fn apply(&mut self, mut p: usize, f: F::F) {
+    pub fn apply(&mut self, p: usize, f: F::F) {
         assert!(p < self.n);
-        p += self.size;
         for i in (1..=self.log).rev() {
-            self.push(p >> i);
+            self.push((p + self.size) >> i);
         }
         self.d[p] = F::mapping(&f, &self.d[p]);
-        for i in 1..=self.log {
-            self.update(p >> i);
-        }
     }
     pub fn apply_range<R>(&mut self, range: R, f: F::F)
     where
@@ -178,8 +108,6 @@
         }
 
         {
-            let l2 = l;
-            let r2 = r;
             while l < r {
                 if l & 1 != 0 {
                     self.all_apply(l, f.clone());
@@ -192,126 +120,29 @@
                 l >>= 1;
                 r >>= 1;
             }
-            l = l2;
-            r = r2;
-        }
-
-        for i in 1..=self.log {
-            if ((l >> i) << i) != l {
-                self.update(l >> i);
-            }
-            if ((r >> i) << i) != r {
-                self.update((r - 1) >> i);
-            }
         }
     }
-
-    pub fn max_right<G>(&mut self, mut l: usize, g: G) -> usize
-    where
-        G: Fn(<F::M as Monoid>::S) -> bool,
-    {
-        assert!(l <= self.n);
-        assert!(g(F::identity_element()));
-        if l == self.n {
-            return self.n;
-        }
-        l += self.size;
-        for i in (1..=self.log).rev() {
-            self.push(l >> i);
-        }
-        let mut sm = F::identity_element();
-        while {
-            // do
-            while l % 2 == 0 {
-                l >>= 1;
-            }
-            if !g(F::binary_operation(&sm, &self.d[l])) {
-                while l < self.size {
-                    self.push(l);
-                    l *= 2;
-                    let res = F::binary_operation(&sm, &self.d[l]);
-                    if g(res.clone()) {
-                        sm = res;
-                        l += 1;
-                    }
-                }
-                return l - self.size;
-            }
-            sm = F::binary_operation(&sm, &self.d[l]);
-            l += 1;
-            //while
-            {
-                let l = l as isize;
-                (l & -l) != l
-            }
-        } {}
-        self.n
-    }
-
-    pub fn min_left<G>(&mut self, mut r: usize, g: G) -> usize
-    where
-        G: Fn(<F::M as Monoid>::S) -> bool,
-    {
-        assert!(r <= self.n);
-        assert!(g(F::identity_element()));
-        if r == 0 {
-            return 0;
-        }
-        r += self.size;
-        for i in (1..=self.log).rev() {
-            self.push((r - 1) >> i);
-        }
-        let mut sm = F::identity_element();
-        while {
-            // do
-            r -= 1;
-            while r > 1 && r % 2 != 0 {
-                r >>= 1;
-            }
-            if !g(F::binary_operation(&self.d[r], &sm)) {
-                while r < self.size {
-                    self.push(r);
-                    r = 2 * r + 1;
-                    let res = F::binary_operation(&self.d[r], &sm);
-                    if g(res.clone()) {
-                        sm = res;
-                        r -= 1;
-                    }
-                }
-                return r + 1 - self.size;
-            }
-            sm = F::binary_operation(&self.d[r], &sm);
-            // while
-            {
-                let r = r as isize;
-                (r & -r) != r
-            }
-        } {}
-        0
-    }
 }
 
-pub struct LazySegtree<F>
+pub struct DualSegtree<F>
 where
     F: MapMonoid,
 {
     n: usize,
     size: usize,
     log: usize,
-    d: Vec<<F::M as Monoid>::S>,
+    d: Vec<F::S>,
     lz: Vec<F::F>,
 }
-impl<F> LazySegtree<F>
+impl<F> DualSegtree<F>
 where
     F: MapMonoid,
 {
-    fn update(&mut self, k: usize) {
-        self.d[k] = F::binary_operation(&self.d[2 * k], &self.d[2 * k + 1]);
-    }
     fn all_apply(&mut self, k: usize, f: F::F) {
-        self.d[k] = F::mapping(&f, &self.d[k]);
         if k < self.size {
             self.lz[k] = F::composition(&f, &self.lz[k]);
+        } else {
+            self.d[k - self.size] = F::mapping(&f, &self.d[k - self.size]);
         }
     }
     fn push(&mut self, k: usize) {
@@ -321,45 +152,16 @@
     }
 }
 
-// TODO is it useful?
-use std::{
-    fmt::{Debug, Error, Formatter, Write},
-    ops::{Bound, RangeBounds},
-};
-impl<F> Debug for LazySegtree<F>
-where
-    F: MapMonoid,
-    F::F: Debug,
-    <F::M as Monoid>::S: Debug,
-{
-    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
-        for i in 0..self.log {
-            for j in 0..1 << i {
-                f.write_fmt(format_args!(
-                    "{:?}[{:?}]\t",
-                    self.d[(1 << i) + j],
-                    self.lz[(1 << i) + j]
-                ))?;
-            }
-            f.write_char('\n')?;
-        }
-        for i in 0..self.size {
-            f.write_fmt(format_args!("{:?}\t", self.d[self.size + i]))?;
-        }
-        Ok(())
-    }
-}
-
 #[cfg(test)]
 mod tests {
-    use std::ops::{Bound::*, RangeBounds};
+    use std::convert::Infallible;
 
-    use crate::{LazySegtree, MapMonoid, Max};
+    use super::{DualSegtree, MapMonoid};
 
-    struct MaxAdd;
-    impl MapMonoid for MaxAdd {
-        type M = Max<i32>;
+    struct RangeAdd(Infallible);
+    impl MapMonoid for RangeAdd {
         type F = i32;
+        type S = i32;
 
         fn identity_map() -> Self::F {
             0
@@ -375,14 +177,15 @@
     }
 
     #[test]
-    fn test_max_add_lazy_segtree() {
+    fn test_range_add_dual_segtree() {
         let base = vec![3, 1, 4, 1, 5, 9, 2, 6, 5, 3];
         let n = base.len();
-        let mut segtree: LazySegtree<MaxAdd> = base.clone().into();
+        let mut segtree: DualSegtree<RangeAdd> = base.clone().into();
         check_segtree(&base, &mut segtree);
 
-        let mut segtree = LazySegtree::<MaxAdd>::new(n);
         let mut internal = vec![i32::min_value(); n];
+        let mut segtree = DualSegtree::<RangeAdd>::from(internal.clone());
+
         for i in 0..n {
             segtree.set(i, base[i]);
             internal[i] = base[i];
@@ -411,69 +214,11 @@
     }
 
     //noinspection DuplicatedCode
-    fn check_segtree(base: &[i32], segtree: &mut LazySegtree<MaxAdd>) {
+    fn check_segtree(base: &[i32], segtree: &mut DualSegtree<RangeAdd>) {
         let n = base.len();
         #[allow(clippy::needless_range_loop)]
         for i in 0..n {
             assert_eq!(segtree.get(i), base[i]);
         }
-
-        check(base, segtree, ..);
-        for i in 0..=n {
-            check(base, segtree, ..i);
-            check(base, segtree, i..);
-            if i < n {
-                check(base, segtree, ..=i);
-            }
-            for j in i..=n {
-                check(base, segtree, i..j);
-                if j < n {
-                    check(base, segtree, i..=j);
-                    check(base, segtree, (Excluded(i), Included(j)));
-                }
-            }
-        }
-        assert_eq!(
-            segtree.all_prod(),
-            base.iter().max().copied().unwrap_or(i32::min_value())
-        );
-        for k in 0..=10 {
-            let f = |x| x < k;
-            for i in 0..=n {
-                assert_eq!(
-                    Some(segtree.max_right(i, f)),
-                    (i..=n)
-                        .filter(|&j| f(base[i..j]
-                            .iter()
-                            .max()
-                            .copied()
-                            .unwrap_or(i32::min_value())))
-                        .max()
-                );
-            }
-            for j in 0..=n {
-                assert_eq!(
-                    Some(segtree.min_left(j, f)),
-                    (0..=j)
-                        .filter(|&i| f(base[i..j]
-                            .iter()
-                            .max()
-                            .copied()
-                            .unwrap_or(i32::min_value())))
-                        .min()
-                );
-            }
-        }
-    }
-
-    fn check(base: &[i32], segtree: &mut LazySegtree<MaxAdd>, range: impl RangeBounds<usize>) {
-        let expected = base
-            .iter()
-            .enumerate()
-            .filter_map(|(i, a)| Some(a).filter(|_| range.contains(&i)))
-            .max()
-            .copied()
-            .unwrap_or(i32::min_value());
-        assert_eq!(segtree.prod(range), expected);
     }
 }

(304行削除、49行追加)

作成した双対セグ木のソースコード

作成した双対セグ木のソースコードは以下の通りです。 このコードは自由に使用していただいて構いませんが、バグがないことは保証しません。

use std::ops::{Bound, RangeBounds};

fn ceil_pow2(n: u32) -> u32 {
    32 - n.saturating_sub(1).leading_zeros()
}

pub trait MapMonoid {
    type F: Clone;
    type S: Clone;
    fn identity_map() -> Self::F;
    fn mapping(f: &Self::F, x: &Self::S) -> Self::S;
    fn composition(f: &Self::F, g: &Self::F) -> Self::F;
}

impl<F: MapMonoid> Default for DualSegtree<F>
where
    F::S: Default,
{
    fn default() -> Self {
        Self::new(0)
    }
}
impl<F: MapMonoid> DualSegtree<F> {
    pub fn new(n: usize) -> Self
    where
        F::S: Default,
    {
        vec![F::S::default(); n].into()
    }
}

impl<F: MapMonoid> From<Vec<F::S>> for DualSegtree<F>
where
    F::S: Default,
{
    fn from(v: Vec<F::S>) -> Self {
        let n = v.len();
        let log = ceil_pow2(n as u32) as usize;
        let size = 1 << log;
        let mut d = vec![F::S::default(); size];
        let lz = vec![F::identity_map(); size];
        d[..n].clone_from_slice(&v);
        DualSegtree {
            n,
            size,
            log,
            d,
            lz,
        }
    }
}

impl<F: MapMonoid> DualSegtree<F> {
    pub fn set(&mut self, p: usize, x: F::S) {
        assert!(p < self.n);
        for i in (1..=self.log).rev() {
            self.push((p + self.size) >> i);
        }
        self.d[p] = x;
    }

    pub fn get(&mut self, p: usize) -> F::S {
        assert!(p < self.n);
        for i in (1..=self.log).rev() {
            self.push((p + self.size) >> i);
        }
        self.d[p].clone()
    }

    pub fn apply(&mut self, p: usize, f: F::F) {
        assert!(p < self.n);
        for i in (1..=self.log).rev() {
            self.push((p + self.size) >> i);
        }
        self.d[p] = F::mapping(&f, &self.d[p]);
    }
    pub fn apply_range<R>(&mut self, range: R, f: F::F)
    where
        R: RangeBounds<usize>,
    {
        let mut r = match range.end_bound() {
            Bound::Included(r) => r + 1,
            Bound::Excluded(r) => *r,
            Bound::Unbounded => self.n,
        };
        let mut l = match range.start_bound() {
            Bound::Included(l) => *l,
            Bound::Excluded(l) => l + 1,
            // TODO: There are another way of optimizing [0..r)
            Bound::Unbounded => 0,
        };

        assert!(l <= r && r <= self.n);
        if l == r {
            return;
        }

        l += self.size;
        r += self.size;

        for i in (1..=self.log).rev() {
            if ((l >> i) << i) != l {
                self.push(l >> i);
            }
            if ((r >> i) << i) != r {
                self.push((r - 1) >> i);
            }
        }

        {
            while l < r {
                if l & 1 != 0 {
                    self.all_apply(l, f.clone());
                    l += 1;
                }
                if r & 1 != 0 {
                    r -= 1;
                    self.all_apply(r, f.clone());
                }
                l >>= 1;
                r >>= 1;
            }
        }
    }
}

pub struct DualSegtree<F>
where
    F: MapMonoid,
{
    n: usize,
    size: usize,
    log: usize,
    d: Vec<F::S>,
    lz: Vec<F::F>,
}
impl<F> DualSegtree<F>
where
    F: MapMonoid,
{
    fn all_apply(&mut self, k: usize, f: F::F) {
        if k < self.size {
            self.lz[k] = F::composition(&f, &self.lz[k]);
        } else {
            self.d[k - self.size] = F::mapping(&f, &self.d[k - self.size]);
        }
    }
    fn push(&mut self, k: usize) {
        self.all_apply(2 * k, self.lz[k].clone());
        self.all_apply(2 * k + 1, self.lz[k].clone());
        self.lz[k] = F::identity_map();
    }
}

#[cfg(test)]
mod tests {
    use std::convert::Infallible;

    use super::{DualSegtree, MapMonoid};

    struct RangeAdd(Infallible);
    impl MapMonoid for RangeAdd {
        type F = i32;
        type S = i32;

        fn identity_map() -> Self::F {
            0
        }

        fn mapping(&f: &i32, &x: &i32) -> i32 {
            f + x
        }

        fn composition(&f: &i32, &g: &i32) -> i32 {
            f + g
        }
    }

    #[test]
    fn test_range_add_dual_segtree() {
        let base = vec![3, 1, 4, 1, 5, 9, 2, 6, 5, 3];
        let n = base.len();
        let mut segtree: DualSegtree<RangeAdd> = base.clone().into();
        check_segtree(&base, &mut segtree);

        let mut internal = vec![i32::min_value(); n];
        let mut segtree = DualSegtree::<RangeAdd>::from(internal.clone());

        for i in 0..n {
            segtree.set(i, base[i]);
            internal[i] = base[i];
            check_segtree(&internal, &mut segtree);
        }

        segtree.set(6, 5);
        internal[6] = 5;
        check_segtree(&internal, &mut segtree);

        segtree.apply(5, 1);
        internal[5] += 1;
        check_segtree(&internal, &mut segtree);

        segtree.set(6, 0);
        internal[6] = 0;
        check_segtree(&internal, &mut segtree);

        segtree.apply_range(3..8, 2);
        internal[3..8].iter_mut().for_each(|e| *e += 2);
        check_segtree(&internal, &mut segtree);

        segtree.apply_range(2..=5, 7);
        internal[2..=5].iter_mut().for_each(|e| *e += 7);
        check_segtree(&internal, &mut segtree);
    }

    //noinspection DuplicatedCode
    fn check_segtree(base: &[i32], segtree: &mut DualSegtree<RangeAdd>) {
        let n = base.len();
        #[allow(clippy::needless_range_loop)]
        for i in 0..n {
            assert_eq!(segtree.get(i), base[i]);
        }
    }
}

テスト

双対セグ木のコードに付属しているテストが通ることは確認しました。また、作成した双対セグ木で ABC332 F - Random Update Query が通ることを確認しました（双対セグ木を使った ABC332 F の提出）。

期待値DPと確率DPの例

こんにちは、ぱるまです。

最近 AtCoder で出た期待値DP/確率DPの問題を整理してみました。具体的には以下の3問を整理してみました。

• ABC326 E - Revenge of "The Salary of AtCoder Inc."
• ABC333 F - Bomb Game 2
• ARC174 C - Catastrophic Roulette

共通の話

• 後ろから考えるとうまく解けがちです。
• 遷移のグラフを作って、（入っていく方の辺より）出ていく方の辺を考えると良さそうです
• グラフのループがある場合は一次方程式が立てられるはずです。
  • 一次方程式を解く際はゼロ除算に注意です。

ABC326 E - Revenge of "The Salary of AtCoder Inc."

問題の概要は以下の通りです。

$N$ 面ダイスと変数 $x=0$ を用意する。
以下の手順を繰り返す。

ダイスを1度振って出た目を $y$ をする。

• もし $x < y$ ならば $A_y$ 円支給し、$x = y$ と更新する。
• そうでないならば終了する。

このとき支給された金額の合計値の期待値 ${}\bmod{998244353}$ は？

[制約]

• $1 \leq N \leq 3 \times 10^5$
• $0 \leq A_i < 998244353$

これは以下のように DP をすることで解けます。（画像では $N$ が $n$ になっていますが気にしないてください）

遷移をグラフにして、dp[i] を考えるときは頂点 i から出ていく辺に着目して、漸化式を作りました。

この問題はループがないシンプルなパターンです。

自分の提出(Rust): https://atcoder.jp/contests/abc326/submissions/51019715

ABC333 F - Bomb Game 2

問題の概要は以下の通りです。

$N$ 人の人が一列に並んでいる。人 $i$ は先頭から $i$ 番目に並んでいる。
以下の操作を $1$ 人になるまで繰り返します。

• 列の先頭にいる人を $1/2$ の確率で列から取り除く。取り除かない場合は列の末尾に移す。

$i=1,2,\ldots, N$ に対して、人 $i$ が列にいる最後の1人になる確率 ${}\bmod{998244353}$ は？

[制約]

• $2 \leq N \leq 3000$

これは以下のように DP をすることで解けます。

ループの部分は一次方程式を解くことで、DP 配列の各値が求められます。

ゼロ除算が発生しないことの確認

${}\bmod{998244353}$ で $1/(2^{i+1}) \equiv 1$ のとき、ゼロ除算が発生します。

$1/(2^{i+1}) \equiv 1$ は $2^{i+1} \equiv 1$ と同値です。

ここで、$2^{k} \equiv 1$ となる最小の正整数 $k$ は $k = (998244353-1)/2$ であることが実験からわかります*1。

# 実験の例 (Python)
x = 1
for i in range(1, 998244353):
    x = (x * 2) % 998244353
    if x == 1:
        print(i) # (998244353-1)/2 = 499122176 が出力される
        break

$1 \leq i \leq N$ で $N=3000$ なので、 $2^{i+1} \not\equiv 1$ が成り立ちます。 よって、今回の制約下ではゼロ除算が発生しません

自分の提出(Rust): https://atcoder.jp/contests/abc333/submissions/48595243

ARC 174 C - Catastrophic Roulette

問題の概要は以下の通りです。

整数 $1,2,\ldots ,N$ が均等な確率で出るルーレットがある。 2人で以下のゲームをする

• 先攻と後攻が交互にルーレットを回す。
  • 出た整数が今までに出ていないものであった場合、何も起こらない。
  • そうでない場合、ルーレットを回したプレイヤーが罰金 $1$ 円を払う。
• $N$ 個の整数すべてが少なくとも $1$ 度出たとき、ただちにゲームが終了する。

先行後攻それぞれについて、ゲームが終了するまでに支払う罰金の期待値 ${}\bmod{998244353}$ は？

[制約]

• $1 \leq N \leq 10^6$

これは以下のように DP をすることで解けます。

ABC333 F - Bomb Game 2 と大体同じです。

ゼロ除算が発生しないことの確認

${}\bmod{998244353}$ で $(i/N)^2 \equiv 1$ のとき、ゼロ除算が発生します。

$(i/N)^2 \equiv 1$ は $i^2 \equiv N^2$ つまり $(N + i)(N-i) \equiv 0$ と同値です。 ここで、$0 \leq i < N$ と問題の制約 $2 \leq N \leq 3000$ から $N{}+i$ と $N-i$ はともに $998244353$ の倍数になりません。 つまり、$(N + i)(N-i) \not \equiv 0$ です。

よって、今回の制約下ではゼロ除算が発生しません。

自分の提出(Rust): https://atcoder.jp/contests/arc174/submissions/52441869

まとめ

確率DP/期待値DP を使う問題を3問紹介しました。どれも遷移のグラフを作って後ろから考えるというよく似た解き方で解けることがわかりました。

*1:ラグランジュの定理から、${}\bmod{998244353}$ で $2^{k} \equiv 1$ となる最小の正整数 $k$ は $998244343-1$ の約数であること言えます。今回の実験では特に考慮せず全探索をしています。