Official

B - ボールの転がり / Rolling Ball Editorial by admin

Gemini 3.0 Flash

概要

\(N\) 個のボールを階段状のマスに転がし、それぞれのボールがどのマスで止まるかを求める問題です。ボールはある段差を越えるために必要な重さを満たさない場合、その直前のマスで停止します。

考察

1. ボールが止まる条件の整理

ボール \(j\) がマス \(k+1\) まで到達するためには、そこに至るまでのすべての段差 \(1, 2, \dots, k\) を越える必要があります。 つまり、以下の条件をすべて満たす必要があります。 - \(B_j \geq W_1\) - \(B_j \geq W_2\) - \(\vdots\) - \(B_j \geq W_k\)

これを言い換えると、「ボールの重さ \(B_j\) が、それまでの段差の重さの最大値以上であること」、すなわち \(B_j \geq \max(W_1, W_2, \dots, W_k)\) が条件となります。

もし、ある \(i\)\(B_j < W_i\) となった場合、ボールはマス \(i\) で停止します。

2. 素朴な解法とその限界

各ボールについて、マス \(1\) から順に「次の段差を越えられるか」を判定していく方法が考えられます。 しかし、この方法では最悪の場合、各ボールについて \(M\) 回の判定を行うことになり、全体の計算量は \(O(N \times M)\) となります。 本問題の制約は \(N, M \leq 2 \times 10^5\) であるため、最大で \(4 \times 10^{10}\) 回程度の計算が必要になり、実行時間制限に間に合いません。

3. 高速化のアイデア:累積最大値と二分探索

「ボールがどこまで進めるか」を効率よく探すために、以下の性質を利用します。 - 段差を越えるために必要な条件は \(B_j \geq \max(W_1, \dots, W_k)\) である。 - \(P_k = \max(W_1, \dots, W_k)\) と定義すると、配列 \(P\)\(P_1 \leq P_2 \leq \dots \leq P_{M-1}\) となり、広義単調増加(値が減ることがない)になる。

配列がソートされた状態(単調増加)であれば、二分探索を用いることで、特定の条件を満たす境界を高速に見つけることができます。 具体的には、二分探索を使って \(P_k \leq B_j\) を満たす最大の \(k\) を探せば、それが「越えられる段差の数」となります。

アルゴリズム

  1. 累積最大値の計算: 段差の重さの配列 \(W\) から、累積最大値の配列 \(P\) を作成します。 \(P_i = \max(W_1, W_2, \dots, W_i)\) とします。
  2. 各ボールに対する二分探索: 各ボールの重さ \(B_j\) について、配列 \(P\) 内で \(B_j\) 以下の値がどこまで続いているかを二分探索(Pythonでは bisect_right など)で求めます。
  3. 停止マスの決定: 越えられた段差の数を \(k\) とすると、ボールが停止するマスの番号は \(k+1\) となります。

具体例

  • 段差の重さ \(W = [3, 5, 2, 8]\)
  • 累積最大値 \(P = [3, 5, 5, 8]\)
  • ボールの重さ \(B_j = 6\) の場合:
    • \(P\) の中で \(6\) 以下の要素は \(3, 5, 5\)\(3\) つです。
    • つまり \(3\) つの段差を越えられるので、停止するマスは \(3 + 1 = 4\) 番目となります。

計算量

  • 時間計算量: \(O(M + N \log M)\)
    • 累積最大値の作成に \(O(M)\) かかります。
    • 各ボール(\(N\) 個)に対して二分探索 \(O(\log M)\) を行うため、合計で \(O(N \log M)\) となります。
  • 空間計算量: \(O(N + M)\)
    • 入力データと累積最大値の配列を保持するために必要なメモリ量です。

実装のポイント

  • 高速な入出力: \(N, M\) が大きいため、Pythonでは sys.stdin.read().split() を使って一括で読み込んだり、sys.stdout.write を使ってまとめて出力したりすることで、入出力のボトルネックを解消できます。

  • 二分探索のライブラリ: bisect モジュールの bisect_right を使うことで、簡潔かつ高速に二分探索を実装できます。

    ソースコード

import sys
from bisect import bisect_right
from itertools import accumulate

def solve():
    # 入力を一括で読み込み、スペースや改行で分割してリストにする
    # 大量のデータを扱う場合、input()を繰り返すより高速です
    try:
        input_data = sys.stdin.read().split()
    except EOFError:
        return
    
    if not input_data:
        return
    
    # N: ボールの数, M: マスの数
    N = int(input_data[0])
    M = int(input_data[1])
    
    # W: 各段差を越えるために必要な重さ (M-1個)
    # インデックス2からMまでの要素がW_1, ..., W_{M-1}に対応します
    W = list(map(int, input_data[2:M+1]))
    
    # B: 各ボールの重さ (N個)
    # インデックスM+1から最後までの要素がB_1, ..., B_Nに対応します
    B = list(map(int, input_data[M+1:M+1+N]))
    
    # P: Wの累積最大値
    # ボールが段差iを越えるためには、それまでのすべての段差(W_1からW_i)を越える必要があります。
    # つまり、ボールの重さ b が max(W_1, ..., W_i) 以上である必要があります。
    # 累積最大値をとることで、配列Pは広義単調増加になり、二分探索が可能になります。
    P = list(accumulate(W, max))
    
    results = []
    for b in B:
        # bisect_rightを用いて、Pの中でb以下の要素がいくつあるかを求めます。
        # これは、そのボールが越えることができる段差の数に相当します。
        # 段差をk個越えた場合、ボールはマスk+1で停止します。
        steps_passed = bisect_right(P, b)
        results.append(str(steps_passed + 1))
    
    # すべての結果を改行で繋いで一度に出力します
    sys.stdout.write('\n'.join(results) + '\n')

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-preview によって生成されました。

posted:
last update: