期待値DPと確率DPの例

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

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

共通の話

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

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$ の約数であること言えます。今回の実験では特に考慮せず全探索をしています。

最大区間和の求め方4通り(DP/モノイド/累積和累積min/Kadane's Algorithm)

ARC174 A - A Multiply で最大区間和を求める問題が出てきました。

コンテスト後、最大区間和を求める方法がいくつかあるのを知りました。 この記事では最大区間和を求める方法を4つ紹介します。

以下、長さ n の整数配列 xs の最大区間和を求める問題を考えます。

定義

長さ n の整数配列 xs に対して、

xs の最大区間和 := max {xs[begin] + … + xs[end-1] | 0 ≤ begin ≤ end ≤ n}
xs の最大 prefix 和 := max {xs[0] + … + xs[end-1] | 0 ≤ end ≤ n}
xs の最大 suffix 和 := max {xs[begin] + … + xs[n-1] | 0 ≤ begin ≤ n}

と定義します。

1. 最大区間和DP

まずは最大区間和をDPで直接求める方法を説明します。

最大区間和のDPと最大 suffix 和のDPを以下で定義します

  • dp_internal_sum[i] = [0, i) における最大区間和
    • これ単独だと漸化式が立たないので、以下の dp_suffix_sum が必要
  • dp_suffix_sum[i] = [0, i) における最大 suffix 和

(i の範囲は 0 ≤ i ≤ n)

DPの初期値は以下の通りです。

  • dp_internal_sum[0] = 0
  • dp_suffix_sum[0] = 0

DPの漸化式は以下の通りです。

  • dp_internal_sum[i+1] = max(dp_internal_sum[i], dp_suffix_sum[i] + xs[i])
    • xs[i] を使うかどうかで場合分け
  • dp_suffix_sum[i+1] = max(xs[i], dp_suffix_sum[i] + xs[i])
    • [0, i) を使うかどうかで場合分け

コードは以下のようになります(Rust)。

fn range_max(xs: &[i64]) -> i64 {
    let n = xs.len()

    // dp_internal_sum[i] = [0, i) での最大区間和
    let mut dp_internal_sum = vec![0; n + 1];
    // dp_suffix_sum[i] = [0, i) での最大 suffix 和
    let mut dp_suffix_sum = vec![0; n + 1]; 

    dp_internal_sum[0] = 0;
    dp_suffix_sum[0] = 0;

    for i in 0..n {
        dp_internal_sum[i + 1] = i64::max(dp_internal_sum[i], dp_suffix_sum[i] + xs[i]);
        dp_suffix_sum[i + 1] = i64::max(xs[i], dp_suffix_sum[i] + xs[i]);
    }
    dp_internal_sum[n]
}

2. 最大区間和モノイド

以下の4つの情報を元に持つモノイドを考えます。

  • 最大区間
  • 最大 prefix 和
  • 最大 suffix 和

モノイドの二項演算は2つの数列を連結させたときの和・最大区間和・最大 prefix 和・最小 prefix 和として定義します

この4つの情報が必要な理由。

欲しい情報は最大区間和です。

2つの数列を連結させたときの最大区間和を計算するには、最大 prefix 和と最大 suffix 和が必要になります

さらに、2つの数列を連結させたときの最大 prefix 和・最大 suffix 和を計算するには和が必要になります

以上から、和・最大区間和・最大 prefix 和・最小 prefix 和が必要となります。

このモノイドは具体的には以下のように定義できます

#[derive(Clone, Debug, PartialEq, Eq)]
struct RangeSumMax {
    prefix_sum_max: i64,
    internal_sum_max: i64,
    suffix_sum_max: i64,
    sum: i64,
}
impl RangeSumMax {
        // 長さ1の数列 [x] に対応する情報
    fn unit(x: i64) -> RangeSumMax {
        RangeSumMax {
            prefix_sum_max: x,
            internal_sum_max: x,
            suffix_sum_max: x,
            sum: x,
        }
    }
}

struct Concat(Infallible);
impl Monoid for Concat {
    type S = RangeSumMax;
    // 空列 [] に対応する情報
    fn identity() -> Self::S {
        RangeSumMax {
            prefix_sum_max: 0,
            internal_sum_max: 0,
            suffix_sum_max: 0,
            sum: 0,
        }
    }

    fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
        RangeSumMax {
            prefix_sum_max: i64::max(a.prefix_sum_max, a.sum + b.prefix_sum_max),
            internal_sum_max: {
                *[
                    a.internal_sum_max,
                    b.internal_sum_max,
                    a.suffix_sum_max + b.prefix_sum_max,
                ]
                .iter()
                .max()
                .unwrap()
            },
            suffix_sum_max: i64::max(b.suffix_sum_max, b.sum + a.suffix_sum_max),
            sum: a.sum + b.sum,
        }
    }
}

モノイドの総積から最大区間和が求まります。

fn range_max(xs: &[i64]) -> i64 {
    let xs_info = xs
        .iter()
        .copied()
        .map(RangeSumMax::unit)
        .fold(Concat::identity(), |acc, x| {
            Concat::binary_operation(&acc, &x)
        });
    xs_info.internal_sum_max;
}

やっていることは最大区間和DPとほとんど同じです。

モノイドを定義するのがつらいですが、セグ木にそのまま応用できるメリットがあります。

3. 累積和の累積min

xs の累積和を cumsum とします。

つまり

  • cumsum[0] = 0
  • cumsum[i] = xs[0] + … + xs[i-1] (0 ≤ i ≤ n)

とします。

区間和は xs[begin] + .. + xs[end-1] = cumsum[end] - cumsum[begin] と表せます。

つまり、最大区間和は

max { cumsum[end] - cumsum[begin] | 0 ≤ begin ≤ end ≤ n}

で求まります。

ここで、end を固定して、

max { cumsum[end] - cumsum[begin] | 0 ≤ begin ≤ end}
= cumsum[end] - min {cumsum[begin] | 0 ≤ begin ≤ end}

を求めて、各 end での値で max を取ればよいです。

min {cumsum[begin] | 0 ≤ begin ≤ end} の部分が累積和の累積 min です。

コードにすると以下のようになります。

use ac_library::Min;

fn range_max(xs: &[i64]) -> i64 {
    let n = xs.len();
    let cumsum = CumSum::new(xs).cumsum;
    let cumsum_cummin = CumMonoid::<Min<i64>>::new(&cumsum);

    (0..=n)
        .map(|end| {
            // cumsum[end] - min {cumsum[begin] | 0 <= begin <= end}
            cumsum[end] - cumsum_cummin.prefix_prod(end + 1)
        })
        .max()
        .unwrap()
}

補足: CumSum と CumMonoid の定義

pub mod cumsum {
    pub struct CumSum {
        pub cumsum: Vec<i64>,
    }
    impl CumSum {
        /// 計算量: O(|xs|)
        pub fn new(xs: &[i64]) -> CumSum {
            let mut cumsum = vec![0; xs.len() + 1];
            for i in 1..xs.len() + 1 {
                cumsum[i] = cumsum[i - 1] + xs[i - 1];
            }
            CumSum { cumsum }
        }
        /// 計算量: O(1)
        pub fn get_interval_sum(&self, begin: usize, end: usize) -> i64 {
            self.cumsum[end] - self.cumsum[begin]
        }
    }
}

pub mod cum_monoid {
    use ac_library::Monoid;
    pub struct CumMonoid<M>
    where
        M: Monoid,
    {
        prefix_prod: Vec<M::S>,
        suffix_prod: Vec<M::S>,
    }
    impl<M> CumMonoid<M>
    where
        M: Monoid,
    {
        pub fn new(xs: &[M::S]) -> CumMonoid<M> {
            let mut prefix_prod = vec![M::identity(); xs.len() + 1];
            let mut suffix_prod = vec![M::identity(); xs.len() + 1];
            for i in 0..xs.len() {
                prefix_prod[i + 1] = M::binary_operation(&prefix_prod[i], &xs[i]);
            }
            for i in (0..xs.len()).rev() {
                suffix_prod[i] = M::binary_operation(&xs[i], &suffix_prod[i + 1]);
            }
            CumMonoid {
                prefix_prod,
                suffix_prod,
            }
        }
        /// [0, i) の総積 (前から累積)
        pub fn prefix_prod(&self, i: usize) -> M::S {
            self.prefix_prod[i].clone()
        }
        /// [i, n) の総積 (後ろから累積)
        pub fn suffix_prod(&self, i: usize) -> M::S {
            self.suffix_prod[i].clone()
        }
        /// [0, i), [i + 1, n) の区間で総積を取る
        pub fn prod_without1(&self, i: usize) -> M::S {
            M::binary_operation(&self.prefix_prod[i], &self.suffix_prod[i + 1])
        }
        pub fn prod_without_range(&self, l: usize, r: usize) -> M::S {
            M::binary_operation(&self.prefix_prod[l], &self.suffix_prod[r])
        }
    }
}

この解法は ARC174 A の公式解法と同じものです。

また、ABC331 E - Set Meal の主菜全探索と同じ考え方です(今回は end を全探索をして最大化しています)。

4. Kadane’s Algorithm

最大区間和は max {xs[begin] + .. + xs[end-1] | 0 ≤ begin ≤ end ≤ n } です。

ここで、end を固定して、

dp[end] = max {xs[begin] + ... + xs[end-1] | 0 ≤ begin ≤ end }

とします。([0,end) での最大 suffix sum)

dp の漸化式は以下のようになります。

  • dp[0] = 0
  • dp[end+1] = max(dp[end] + xs[end], xs[end])

このとき、最大区間和は max {dp[end] | 0 ≤ end ≤ n} と求まります。

fn range_max(xs: &[i64]) -> i64 {
    let n = xs.len();
    let mut dp = vec![0; n + 1];
    dp[0] = 0;

    for end in 0..n {
        dp[end + 1] = i64::max(dp[end] + xs[end], xs[end]);
    }
    *dp.iter().max().unwrap()
}

このアルゴリズムは Kadane’s Algorithm と呼ばれています。

参考: Kadane's Algorithm | 最大部分配列 問題 - Ark's Blog

おわりに

この記事では、最大区間和を求める方法を4つ紹介しました。

  1. 最大区間和DP
  2. 最大区間和モノイド
  3. 累積和の累積min
  4. Kadane’s Algorithm

どの方法の考え方も他の問題で使えそうなので、覚えておきたいです。