公式

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

Gemini 3.0 Flash

Overview

This is a problem where \(N\) balls are rolled across staircase-shaped cells, and we need to determine which cell each ball stops at. A ball stops at the cell just before a step if it doesn’t have enough weight to overcome that step.

Analysis

1. Organizing the Stopping Conditions

For ball \(j\) to reach cell \(k+1\), it must overcome all steps \(1, 2, \dots, k\) along the way. In other words, all of the following conditions must be satisfied: - \(B_j \geq W_1\) - \(B_j \geq W_2\) - \(\vdots\) - \(B_j \geq W_k\)

Rephrasing this, the condition is “the ball’s weight \(B_j\) must be greater than or equal to the maximum weight among all steps encountered so far”, i.e., \(B_j \geq \max(W_1, W_2, \dots, W_k)\).

If at some \(i\) we have \(B_j < W_i\), then the ball stops at cell \(i\).

2. Naive Solution and Its Limitations

One approach is to check for each ball, starting from cell \(1\), whether it can overcome the next step. However, in the worst case, this requires \(M\) checks per ball, resulting in an overall time complexity of \(O(N \times M)\). Given the constraints \(N, M \leq 2 \times 10^5\), this would require up to around \(4 \times 10^{10}\) operations, which exceeds the time limit.

3. Optimization Idea: Prefix Maximum and Binary Search

To efficiently find how far each ball can travel, we use the following property: - The condition to overcome steps is \(B_j \geq \max(W_1, \dots, W_k)\). - Defining \(P_k = \max(W_1, \dots, W_k)\), the array \(P\) satisfies \(P_1 \leq P_2 \leq \dots \leq P_{M-1}\), making it non-decreasing (values never decrease).

When an array is sorted (monotonically non-decreasing), we can use binary search to quickly find the boundary satisfying a specific condition. Specifically, by using binary search to find the largest \(k\) such that \(P_k \leq B_j\), we obtain the “number of steps that can be overcome.”

Algorithm

  1. Computing the Prefix Maximum: From the step weight array \(W\), create the prefix maximum array \(P\). Set \(P_i = \max(W_1, W_2, \dots, W_i)\).
  2. Binary Search for Each Ball: For each ball’s weight \(B_j\), use binary search (e.g., bisect_right in Python) to find how far values in array \(P\) that are less than or equal to \(B_j\) extend.
  3. Determining the Stopping Cell: If the number of steps overcome is \(k\), the cell number where the ball stops is \(k+1\).

Concrete Example

  • Step weights \(W = [3, 5, 2, 8]\)
  • Prefix maximum \(P = [3, 5, 5, 8]\)
  • For a ball with weight \(B_j = 6\):
    • The elements in \(P\) that are \(\leq 6\) are \(3, 5, 5\) — a total of \(3\).
    • This means \(3\) steps can be overcome, so the ball stops at cell \(3 + 1 = 4\).

Complexity

  • Time Complexity: \(O(M + N \log M)\)
    • Computing the prefix maximum takes \(O(M)\).
    • Binary search \(O(\log M)\) is performed for each of the \(N\) balls, totaling \(O(N \log M)\).
  • Space Complexity: \(O(N + M)\)
    • This is the memory needed to store the input data and the prefix maximum array.

Implementation Notes

  • Fast I/O: Since \(N, M\) can be large, in Python you can eliminate I/O bottlenecks by reading all input at once with sys.stdin.read().split() and writing output collectively with sys.stdout.write.

  • Binary Search Library: Using bisect_right from the bisect module allows for a concise and efficient binary search implementation.

    Source Code

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()

This editorial was generated by gemini-3-flash-preview.