公式

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

Qwen3-Coder-480B

Overview

A problem where you simulate how far each ball rolls on a staircase-shaped grid based on its weight.

Analysis

When a ball rolls starting from square \(1\), it stops at the square just before “the first step whose weight exceeds the ball’s weight \(B_j\).” In other words, the position where the ball stops corresponds to “the first value among \(W_1, W_2, ..., W_{M-1}\) that is greater than \(B_j\), scanning from left to right.”

However, if we compare from the beginning each time, it takes \(O(M)\) in the worst case, resulting in \(O(NM)\) overall, which is too slow (since \(N, M \sim 2 \times 10^5\), up to \(4 \times 10^{10}\) operations would be needed). Therefore, we need to think about how to efficiently determine how far each ball can advance.

A key observation is that whether a ball can reach square \(i\) can be determined by comparing the ball’s weight with the maximum value among all steps \(W_1\) through \(W_{i-1}\) up to that square. In other words, if we precompute “the maximum value from \(W_1\) to \(W_{i-1}\)” for each \(i\), we can use binary search to find the stopping position for each ball.

By preprocessing the prefix maximums and performing binary search, we can efficiently compute the answer.

Algorithm

  1. Read the step weights \(W_1, ..., W_{M-1}\) and ball weights \(B_1, ..., B_N\) from the input.
  2. Prepare an array max_w, where max_w[i] stores the maximum value from \(W_1\) to \(W_i\) (prefix maximum).
  3. For each ball weight \(B_j\), use binary search on max_w to find “the largest index where the value is at most \(B_j\).”
  4. That index + 2 (accounting for 1-based indexing and the fact that the next square is the stopping position) gives the answer.

For example: - When \(W = [3, 1, 4]\) - max_w = [3, 3, 4] - If the ball’s weight is \(5\), then bisect.bisect_right(max_w, 5) → 3 → stops at square 4

Complexity

  • Time complexity: \(O(M + N \log M)\)
    • \(O(M)\) for precomputing the prefix maximums
    • \(O(N \log M)\) for binary search on each ball
  • Space complexity: \(O(M + N)\)
    • Used for the input data and the prefix maximum array

Implementation Notes

  • bisect.bisect_right returns “the first position that exceeds the value,” so appropriate index conversion is necessary.

  • sys.stdin.read is used for fast input reading.

  • Results are accumulated in a list and output all at once at the end for faster I/O.

    Source Code

import bisect
import sys

def main():
    input = sys.stdin.read
    data = input().split()
    
    N = int(data[0])
    M = int(data[1])
    
    W = list(map(int, data[2:M+1]))
    B = list(map(int, data[M+1:M+1+N]))
    
    # 各ボールがどこまで転がれるかを計算
    # Wの累積最大値を前計算しておくと効率的
    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])
    
    result = []
    for b in B:
        # max_wの中で b以上となる最大のindex+1 (0-based -> 1-based)
        pos = bisect.bisect_right(max_w, b)
        result.append(str(pos + 1))
    
    print('\n'.join(result))

if __name__ == "__main__":
    main()

This editorial was generated by qwen3-coder-480b.