概要

階段状のマスの上でボールを転がしたとき、各ボールがどこで停止するかを求める問題です。二分探索を用いることで高速に解くことができます。

考察

問題の理解

ボールはマス1からスタートし、自分の重さ \(B_j\) が段差の閾値 \(W_i\) 以上である限り、次のマスへ転がり続けます。

素朴なアプローチとその問題点

各ボールについて、マス1から順番にシミュレーションすると、1つのボールあたり最大 \(O(M)\) 回のチェックが必要です。ボールが \(N\) 個あるので、全体で \(O(NM)\) となり、\(N, M\) が最大 \(2 \times 10^5\) のとき、最大 \(4 \times 10^{10}\) 回の操作となりTLEします。

重要な気づき

ボールがマス \(k\) まで到達するためには、途中のすべての段差を越える必要があります。つまり、マス1からマス \(k\) まで進むには： $\(B_j \geq W_1 \text{ かつ } B_j \geq W_2 \text{ かつ } \cdots \text{ かつ } B_j \geq W_{k-1}\)$

これは以下と同値です： $\(B_j \geq \max(W_1, W_2, \ldots, W_{k-1})\)$

累積最大値の利用

\(\text{max\_W}[i] = \max(W_1, W_2, \ldots, W_{i+1})\) と定義すると、この配列は単調非減少になります（最大値は減ることがないため）。

例：\(W = [3, 1, 4, 2]\) のとき、\(\text{max\_W} = [3, 3, 4, 4]\)

単調非減少な配列に対しては、二分探索が使えます！

アルゴリズム

1. 前処理: 累積最大値配列 \(\text{max\_W}\) を計算する

   • \(\text{max\_W}[0] = W[0]\)
   • \(\text{max\_W}[i] = \max(\text{max\_W}[i-1], W[i])\)（\(i \geq 1\)）

2. 各ボールについて二分探索:

   • ボールの重さ \(b\) に対して、\(\text{max\_W}[i] \leq b\) を満たす最大の \(i\) を求める
   • bisect_right(max_W, b) は「\(b\) より大きい値が最初に現れる位置」を返す
   • よって bisect_right(max_W, b) 個の段差を越えられる
   • 停止するマスは bisect_right(max_W, b) + 1

具体例

\(M = 5\), \(W = [3, 1, 4, 2]\), ボールの重さ \(b = 3\) の場合： - \(\text{max\_W} = [3, 3, 4, 4]\) - bisect_right([3, 3, 4, 4], 3) = 2（インデックス0, 1の値が3以下） - 2個の段差を越えられるので、マス \(2 + 1 = 3\) で停止

計算量

• 時間計算量: \(O(M + N \log M)\)
  • 累積最大値の計算: \(O(M)\)
  • 各ボールに対する二分探索: \(O(\log M)\) × \(N\) 個 = \(O(N \log M)\)
• 空間計算量: \(O(M + N)\)
  • 累積最大値配列: \(O(M)\)
  • 結果の格納: \(O(N)\)

実装のポイント

• bisect_right は「挿入位置」を返すため、そのまま「越えられる段差の数」として使える

• 累積最大値配列は0-indexedで管理し、\(\text{max\_W}[i]\) は \(W[0]\) から \(W[i]\) までの最大値を表す

• 出力が \(N\) 行あるため、1行ずつ print するより '\n'.join() でまとめて出力する方が高速

  ソースコード

import bisect

def main():
    import sys
    input = sys.stdin.readline
    
    N, M = map(int, input().split())
    W = list(map(int, input().split()))
    B = list(map(int, input().split()))
    
    # max_W[i] = max(W[0], W[1], ..., W[i]) を計算
    # これは、マス1からマスi+2まで進むために必要な最小重さ
    max_W = [0] * (M - 1)
    max_W[0] = W[0]
    for i in range(1, M - 1):
        max_W[i] = max(max_W[i - 1], W[i])
    
    # max_W は単調非減少になっている
    # ボールの重さがb のとき、b >= max_W[i] を満たす最大のiを見つける
    # そのボールはマス i+2 まで進める（i+1個の段差を越える）
    
    results = []
    for b in B:
        # b >= max_W[i] を満たす最大の i を二分探索で見つける
        # max_W[i] <= b となる最大の i
        # bisect_right で b より大きい最初の位置を見つけ、その手前まで進める
        idx = bisect.bisect_right(max_W, b)
        # idx個の段差を越えられる -> マス idx+1 に到達
        results.append(idx + 1)
    
    print('\n'.join(map(str, results)))

if __name__ == "__main__":
    main()

---

この解説は claude4.5opus によって生成されました。