公式

E - ボールの転送 / Ball Transfer 解説 by admin 印刷ページ

Qwen3-Coder-480B

Overview

Each of \(N\) people throws a ball to the right, and it is caught by a person satisfying certain conditions. We need to find the total cost, where the cost is the minimum height in the interval from the thrower to the catcher.

Analysis

In this problem, for each person \(i\), the ball reaches “the leftmost person to the right who is taller than them.” That is, for each \(i\), we need to find: $\( j = \min\{k \mid k > i,\ H_k > H_i\} \)$

What about a naive approach?

A naive approach would scan to the right for each \(i\) to find \(j\) satisfying the condition, which takes \(O(N^2)\) in the worst case. Given the constraint \(N \leq 2 \times 10^5\), this would result in TLE.

Improvement: Monotonic Stack

The problem of finding the first larger value to the right of a given element is a classic “Next Greater Element” problem, and the effective way to solve it efficiently is using a Monotonic Stack. The stack holds indices and is maintained so that the heights are in strictly decreasing order. When a new element arrives, we pop indices with smaller heights while updating them.

This allows us to precompute the position of the next taller person next_greater[i] for each element in \(O(N)\).

We need the interval minimum

Next, we need to find the minimum height in the interval from \(i\) to \(j\). This is nothing other than a Range Minimum Query (RMQ), and a Segment Tree is a suitable efficient data structure for this. By building it in advance, each query can be processed in \(O(\log N)\).

Algorithm

  1. For each person \(i\), find the index \(j\) of the first person to the right who is taller than \(H_i\) (using a monotonic stack).
  2. For each pair \((i, j)\), retrieve the minimum value (= cost) in the interval \([i, j]\) using a segment tree.
  3. Sum up all costs and output the result.

Complexity

  • Time complexity: \(O(N \log N)\)
    • Building next_greater with the monotonic stack: \(O(N)\)
    • Building the segment tree: \(O(N)\)
    • Each query: \(O(\log N)\) × \(N\) times → \(O(N \log N)\)
  • Space complexity: \(O(N)\)
    • Both the stack and segment tree are linear in size

Implementation Notes

  • In the monotonic stack, maintain the stack so that heights are in strictly decreasing order.

  • Segment tree range queries often use half-open intervals \([l, r)\), so be careful to pass j+1 rather than r+1 as the endpoint.

  • Initialize the minimum value with a sufficiently large value (e.g., float('inf')).

    Source Code

import sys
from collections import deque

def main():
    input = sys.stdin.read
    data = input().split()
    N = int(data[0])
    H = list(map(int, data[1:]))

    # 各iについて、iより右でH[i]より大きい最小のjを求める
    # これは単調スタックを使って効率的に求められる
    next_greater = [-1] * N
    stack = []  # 単調減少スタック(インデックスを保持)

    for i in range(N):
        while stack and H[stack[-1]] < H[i]:
            idx = stack.pop()
            next_greater[idx] = i
        stack.append(i)

    # 区間の最小値をセグメント木で管理
    class SegmentTree:
        def __init__(self, n, initial_value):
            self.size = 1
            while self.size < n:
                self.size *= 2
            self.tree = [initial_value] * (2 * self.size)
        
        def build(self, arr):
            for i in range(len(arr)):
                self.tree[self.size + i] = arr[i]
            for i in range(self.size - 1, 0, -1):
                self.tree[i] = min(self.tree[2*i], self.tree[2*i+1])
        
        def query(self, l, r):
            # [l, r) の最小値
            l += self.size
            r += self.size
            res = float('inf')
            while l < r:
                if l & 1:
                    res = min(res, self.tree[l])
                    l += 1
                if r & 1:
                    r -= 1
                    res = min(res, self.tree[r])
                l >>= 1
                r >>= 1
            return res

    seg_tree = SegmentTree(N, float('inf'))
    seg_tree.build(H)

    total_cost = 0
    for i in range(N):
        j = next_greater[i]
        if j != -1:
            cost = seg_tree.query(i, j+1)
            total_cost += cost
    
    print(total_cost)

if __name__ == "__main__":
    main()

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

投稿日時:
最終更新: