概要

\(N\) 人の人それぞれがボールを右に投げ、条件を満たす相手にキャッチされるときのコストの合計を求めます。コストは投げた人からキャッチした人までの区間の身長の最小値です。

考察

この問題では、各人 \(i\) について「自分より右にいる、かつ自分より身長が高い人の中で最も左の人」にボールが届くという条件があります。つまり、各 \(i\) に対して、
$\( j = \min\{k \mid k > i,\ H_k > H_i\} \)$
を求めることになります。

素朴な方法では？

ナイーブな方法では、各 \(i\) に対して右側を順番に見て条件を満たす \(j\) を探すため、最悪ケースで \(O(N^2)\) かかります。制約が \(N \leq 2 \times 10^5\) なので、これはTLEになります。

改善策：単調スタック

ある要素より右で初めて大きい値を探す、という問題は典型的な「Next Greater Element」問題であり、これを効率よく解くには単調スタック（Monotonic Stack）を使うのが有効です。スタックにはインデックスを保持し、身長が単調減少になるように管理します。新しい要素が来たら、それより小さい要素を持つインデックスをポップしながら更新していきます。

これにより、各要素に対して次の高身長者の位置 next_greater[i] を前処理 \(O(N)\) で求めることができます。

区間の最小値が必要

次に、見つけた \(i\) から \(j\) までの区間の最小身長を求める必要があります。これは区間最小値クエリ（RMQ）に他ならず、効率的なデータ構造としてセグメント木が適しています。事前に構築することで、各クエリを \(O(\log N)\) で処理できます。

アルゴリズム

1. 各人 \(i\) について、右側で初めて \(H_i\) よりも高い人のインデックス \(j\) を求める（単調スタック使用）。
2. 各ペア \((i, j)\) に対して、区間 \([i, j]\) の最小値（= コスト）をセグメント木で取得。
3. 全てのコストを合計して出力。

計算量

• 時間計算量: \(O(N \log N)\)
  • 単調スタックによる next_greater の構築: \(O(N)\)
  • セグメント木の構築: \(O(N)\)
  • 各クエリ: \(O(\log N)\) × \(N\) 回 → \(O(N \log N)\)
• 空間計算量: \(O(N)\)
  • スタック、セグメント木ともに線形サイズ

実装のポイント

• 単調スタックでは、身長が狭義単調減少になるようにスタックを管理する。

• セグメント木の区間クエリは半開区間 \([l, r)\) で扱うことが多いので、終端 r+1 ではなく j+1 を渡すことに注意。

• 最小値の初期値は十分大きな値（例: float('inf')）にしておく。

  ソースコード

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

この解説は qwen3-coder-480b によって生成されました。