D - Earthquakes Editorial
by 
square1001
この問題を解く戦略
電柱を並べて描いてみましょう。隣り合う電柱が距離 \(H\) より離れていたら、一方が倒れても他方に影響しません。なので、隣り合う電柱の距離が \(H\) 以下の「グループ」ごとに独立に問題を考えることができそうです。
各グループでは、単純化すると、以下の問題を解くことになります。
部分問題 \(m\) 個の電柱がある。左から \(1, \dots, m\) 番目の電柱はそれぞれ \(p_1, \dots, p_m\) 回目の地震で倒れる。隣り合う電柱は衝突しうる。最後に倒れる電柱が左から \(1, \dots, m\) 番目の電柱となる確率をそれぞれ求めよ。（その他の設定は元の問題と同じ ／ \(p_1, \dots, p_m\) はすべて異なる）
以下の戦略で問題を解いていきましょう。
- 部分問題 を解く方法を考える（ステップ 1-4）
- グループごとの結果を効率的にまとめ上げる方法を考える（ステップ 5）
ステップ 1: 電柱の倒れ方の観察
部分問題 で、どのように電柱が倒れていくかを観察すると、どの場面でも、下図のような感じで、左から順に「電柱が左に倒れている区間」「電柱が立っている区間」「電柱が右に倒れている区間」となっていることが分かります（初期状態や最終状態など、区間の電柱の個数が \(0\) 個の場合もあり得ます）。
理由は、上図の状態から電柱が左に倒れた場合でも右に倒れた場合でも、下図のように同様の形で表せる状態になるからです。
ステップ 2: 条件を考えよう
「〇〇の条件を満たすケース」を数え上げたり確率を求めたりする問題の多くは、この条件をもっと考えやすい形にするのが最初のステップになることが多いです。ここでは、部分問題 において最後に倒れる電柱が左から \(i\) 番目の電柱となる条件を考えてみましょう。
まず、\(p_i\) 回目の地震の直前の時点では左から \(i\) 番目の電柱が残っている必要があります。そのためには、その前の地震において
- 左から \(i\) 番目より左の電柱が倒れる場合は左に倒れる必要がある
- 左から \(i\) 番目より右の電柱が倒れる場合は右に倒れる必要がある
理由は、ステップ 1 の事実より、もし左から \(i\) 番目より左の電柱が右に倒れた場合は連鎖的に電柱が倒れていき左から \(i\) 番目の電柱も倒れてしまうからです。左から \(i\) 番目より右の電柱が左に倒れた場合も同様です。
逆に、この 2 つの条件さえ満たされれば左から \(i\) 番目の電柱は倒れようがありません。
その前提で、左から \(i\) 番目の電柱が倒れることによってすべての電柱が倒れた状態になる必要があります。この条件は以下のようになります。
- 左に倒れた場合: 左から \(i+1\) 番目の電柱が立ち残っていなければいいので、\(p_i > p_{i+1}\) または \(i = m\) なら全部倒れた状態になる
- 右に倒れた場合: 左から \(i-1\) 番目の電柱が立ち残っていなければいいので、\(p_i > p_{i-1}\) または \(i = 1\) なら全部倒れた状態になる
ステップ 3: 確率はどのくらい？
3.1: 具体例で考える
まずは以下の例で、最後に倒れる電柱が赤の電柱になる確率を、ステップ 2 の条件をもとに考えてみましょう。
青の電柱は倒れるなら左に、赤の電柱は倒れるなら右に倒れなければなりません。この条件で、\(64\) 回目の地震の直前までシミュレーションすると、以下のようになります。
この 5 つのイベントのすべてで電柱が「正しい方向」（それぞれ右、左、右、右、左）に倒れる必要があります。この確率は \(\left(\frac{1}{2}\right)^5 = \frac{1}{32}\) です。これが \(64\) 回目の地震の直前に赤の電柱が立っている確率です。
ここで \(64\) 回目の地震が起こります。\(64 > 33\) なので、赤の電柱が左に倒れると全部倒れた状態になります。また \(64 > 62\) なので、赤の電柱が右に倒れても全部倒れた状態になります。つまり、今回の例ではどちらに倒れてもよいです。よって、最後に倒れる電柱が赤の電柱になる確率は \(\frac{1}{32}\) です。
3.2: 確率の計算方法
一般化すると、最後に倒れる電柱が左から \(i\) 番目となる確率は
- \(a\) を「左から \(i\) 番目より左の電柱は地震で左に倒れる、左から \(i\) 番目より右の電柱は地震で右に倒れる」として \(p_i\) 回目の地震の直前までシミュレーションしたときに、実際に電柱が倒れる地震の回数
- \(b\) を「\(p_i > p_{i+1}\) または \(i = m\)」「\(p_i > p_{i-1}\) または \(i = 0\)」の 2 つのうちいくつ満たされるか
として、\(\left(\frac{1}{2}\right)^a \times \frac{b}{2}\) と求まります。3.1 節の例では \(a = 5, b = 2\) です。
ここで、\(a\) の値は実際にシミュレーションしなくても求められます。
- シミュレーションで \(p_j \ (j < i)\) 回目の地震で実際に電柱が倒れる条件は \(p_j < \min(p_{j+1}, \dots, p_i)\) です。理由は、電柱 \(j+1, \dots, i-1\) が先に倒れず、かつ \(p_i\) 回目の地震より前である必要があるからです。
- 言い換えれば、電柱 \(i\) より左の電柱が倒れる地震の回数は、「電柱 \(i\) から左に進んだときに \(p_j\) の最小値を更新する回数」です。
- 同様に、電柱 \(i\) より右の電柱が倒れる地震の回数は、「電柱 \(i\) から右に進んだときに \(p_j\) の最小値を更新する回数」です。
- この 2 つの値を足し算することで \(a\) の値が求まります。
ステップ 4: 「最小値の更新回数」の求め方
ステップ 3 の結果より、以下の問題が計算量 \(O(m)\) で解ければ 部分問題 を計算量 \(O(m)\) で解くことができます。
典型問題 長さ \(n\) の配列 \(a_1, \dots, a_n\) が与えられる。各 \(i = 1, \dots, n\) に対して、\(a_i\) から左に進んだときに最小値を更新する回数（つまり \(a_j < \min(a_{j+1}, \dots, a_i)\) となる \(j \ (1 \leq j \leq i-1)\) の個数）を求めよ。
これは実際に AtCoder 400-450 点レベルの典型的な問題ですので、見たことがある方も多いと思います。ここでも解法を紹介します。
典型問題 の解法
現時点での「最小値を更新する場所のリスト」をスタック $S$ で管理して、$i = 1, \dots, n$ の順に答えを求める戦略で解きます。
- $i = 1, \dots, n$ の順に以下を行う
- スタック $S$ が空でなく $a_{(S \text{の一番上の要素})} \geq a_i$ である間、$S$ の一番上の要素の削除を続ける
- $S$ の一番上に $a_i$ を追加する
- $a_i$ から左に進んだときに最小値を更新する回数は、(この時点での $S$ のサイズ) $-$ $1$
ループ $1$ 回の計算量が $O(1)$ になるとは限りませんが、全体の計算量は $O(n)$ になります。理由は、$S$ への追加が起こる回数が $n$ 回なので、削除回数の合計も $n$ 回以下になるからです。
書籍「競技プログラミングの鉄則」 の 8.10 節にも同様のテクニックが解説されていますので、持っている方は確認しておきましょう。
ステップ 5: グループごとの結果を効率的にまとめ上げる
最後に、各グループに対する 部分問題 の結果をまとめて、元の問題の答えを出す方法を考えます。
ちょうど \(t\) 回目の地震ですべての電柱が倒れた状態になる条件は
- 電柱 \(t\) が属するグループで、最後に倒れるのが電柱 \(t\) である
- それ以外の各グループに対して、以下の条件を満たす
- 最後に倒れる電柱が電柱 \(1, \dots, t-1\) のいずれかである
であり、それが起こる確率は各グループごとの 1. / 2. の条件を満たす確率の掛け算となります。これは、部分問題 の結果を直接使って求まります。
問題は、普通の方法では元の問題を解くのに計算量 \(O(N^2)\) かかってしまうということです。なぜなら、最後に倒れる電柱が電柱 \(1, \dots, t-1\) のいずれかである確率を求めるのに各電柱ごとの確率の足し算をする必要があるし、全体の確率も各グループごとの確率の掛け算をする必要があるからです。ここで、\(t = 1, \dots, N\) の順に答えを求めることを利用した工夫で、大量の足し算を回避することができます。
グループの個数を \(K\) とし、電柱 \(1, \dots, N\) が属するグループの番号を \(g_1, \dots, g_N \ (1 \leq g_i \leq K)\) とする。
- 「現時点でグループ \(j\) の電柱がすべて倒れている確率」を \(s_j\) とする。最初 \(s_1, \dots, s_K\) を \(0\) に設定する。
- \(i = 1, \dots, N\) の順に以下を行う。
- グループ \(g_i\) の 部分問題 において、最後に倒れるのが電柱 \(i\) である確率を \(x\) とする。
- \(s_1 \times \dots \times s_{g_i-1} \times x \times s_{g_i+1} \times \dots \times s_K\) を求める。これが \(t = i\) での答えである。
- \(s_{g_i}\) に \(x\) を加算する。
あとは、2. の掛け算さえ効率的に求められればよいです。以下の方法が考えられます。
方法 1: セグメント木を使う
セグメント木を使うと「配列の指定した要素を更新する」「配列上の指定した範囲の要素の掛け算を求める」の 2 種類のクエリを計算量 $O(\log N)$ で処理できます。よって、この問題は計算量 $O(N \log N)$ で解けます。
方法 2: $s_1 \times \dots \times s_K$ の値を管理する
$Z = s_1 \times \dots \times s_K$ の値をリアルタイムで管理することを考えます。すると、答えが $Z \times \frac{x}{s_{g_i}}$ と簡単な式で計算でき、$Z$ の値の更新についても、$s_{g_i}$ に $x$ を加算する際に $Z$ を $\frac{s_{g_i}+x}{s_{g_i}}$ 倍すればよいので簡単です。この問題では答えを素数 $P = 998244353$ で割った余りを求めるので、mod の逆元を使って計算量 $O(\log P)$ で処理することになります。
しかし、$s_{g_i} = 0$ の場合や、$998244353$ の倍数である場合があるので、簡単にはいきません。ひとつの解決策は（$\mathrm{mod} \ 998244353$ で）「$s_1, \dots, s_K$ のうち $0$ 以外のものの掛け算」および「$s_1, \dots, s_K$ のうち $0$ の個数」を管理することで、正しい $Z$ を計算することです。すると、この問題は計算量 $O(N \log P)$ で解けます。
ステップ 6: 実装に関する補足
実際にプログラムを実装するときは、すべてを整数の計算で行うために、確率を \(2^{(グループの電柱の個数)}\) 倍した値を管理して計算するか、mod の逆元を使って計算するかして行いましょう。実装例 では前者の方法を採用しています。
実装例
C++ での実装例は https://atcoder.jp/contests/arc177/submissions/53427881 です。
Python では以下のように実装できます。
MOD = 998244353
# セグメント木
class segment_tree:
def __init__(self, n: int):
self.size = 1 << (n-1).bit_length()
self.v = [0] * (self.size*2)
def add(self, pos: int, x: int):
pos += self.size
self.v[pos] = (self.v[pos] + x) % MOD
while pos > 1:
pos >>= 1
self.v[pos] = self.v[pos * 2 + 0] * self.v[pos * 2 + 1] % MOD
def query(self, l: int, r: int):
l += self.size
r += self.size
res = 1
while l != r:
if l & 1:
res = res * self.v[l] % MOD
l += 1
if r & 1:
r -= 1
res = res * self.v[r] % MOD
l >>= 1
r >>= 1
return res
# 部分問題を解く関数
def calc(n, p):
level = [n-1] * n
sl = list()
for i in range(n):
while sl and sl[-1] > p[i]:
sl.pop()
level[i] -= len(sl)
sl.append(p[i])
sr = list()
for i in range(n-1, -1, -1):
while sr and sr[-1] > p[i]:
sr.pop()
level[i] -= len(sr)
sr.append(p[i])
power2 = [0] * (n+1)
power2[0] = 1
for i in range(n):
power2[i+1] = power2[i] * 2 % MOD
res = [0] * n
for i in range(n):
c = int(i == n-1 or p[i] > p[i+1]) + int(i == 0 or p[i] > p[i-1])
res[i] = power2[level[i]] * c % MOD
return res
# メインの問題を解く関数
def solve(n: int, h: int, x: list):
# 問題をグループごとに分解する
pillars = [(x[i], i) for i in range(n)]
pillars.sort()
p = [pillars[i][1] for i in range(n)]
v = list()
g = list()
pre = 0
m = 0
for i in range(1, n+1):
if i == n or pillars[i][0] - pillars[i-1][0] > h:
v += calc(i-pre, p[pre:i])
g += [m] * (i-pre)
pre = i
m += 1
# グループごとの答えをまとめ上げる
pinv = [-1] * n
for i in range(n):
pinv[p[i]] = i
seg = segment_tree(m)
res = [0]
for i in pinv:
seg.add(g[i], v[i])
res.append(seg.query(0, m))
ans = [(res[i+1]-res[i]) % MOD for i in range(n)]
return ans
# 入力・出力
n, h = map(int, input().split())
x = list(map(int, input().split()))
ans = solve(n, h, x)
print(*ans)
posted:
last update: