Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 263 E問題 Sugoroku 3

AtCoder ABC ABC263 E問題 500 pts

問題

提出解答

問題の概要

マス $1$ からマス $N$ からなるマスがある. 最初, マス $1$ にいる.

マス $1$ からマス $(N-1)$ にはサイコロが1個あり, マス $i$ にあるサイコロには $0$ 以上 $A_i$ 以下の整数が等確率で出る (毎回独立).

マス $N$ に到達するまでサイコロを振り, 出た目だけマスを進むことにする. サイコロを振る回数の期待値を求めよ.

制約

$2 \leq N \leq 2 \times 10^5$
$1 \leq A_i \leq N-i$

解法

$i=1,2, \dots, N$ に対して, $E_i$ でマス $i$ から始めた場合にマス $N$ に到達するまでのサイコロを振る回数の期待値とする.

まず, $E_N=0$ である. 次に, $i=1,2,\dots, N-1$ に対して, マス $i$ から進めるのはマス $i,i+1, \dots, i+A_i$ の $(A_i+1)$ 個であり, それぞれ $\dfrac{1}{A_i+1}$ の確率であるから,

$\displaystyle E_i=\dfrac{1}{A_i+1} \sum_{k=0}^{A_i} E_{i+k}+1$

である. この式を整理すると,

$\displaystyle E_i=\dfrac{1}{A_i} \sum_{k=1}^{A_i} E_{i+k}+\left(1+\dfrac{1}{A_i} \right)$

となる.

ここで, $E_i$ を求めるときには $E_{i+1}, \dots, E_{i+A_i}$ の和という区間和を求めることになる. このような区間和には Binary Indexed Tree を利用することにより, 各 $i$ に対して $\displaystyle \sum_{k=1}^{A_i} E_{i+k}$ を $O(\log N)$ で求めることができる.

従って, 適切な前計算の元で, $E_1$ を $O(N \log N)$ 時間で求めることができる.

また, $E_i$ を求めた後, それ以降は $E_i$ が変わることはないということから, 累積和を利用することにより, $O(N)$ 時間でも求めることができる.