Pythonで海が見える建物を見つけるプログラム|単調スタックによるO(n)解法

問題の概要

さまざまな高さの建物のリストが与えられたとします。ある建物 heights[i] から海が見えるのは、その右側にあるすべての建物がその建物よりも低い場合です。この条件を満たす建物のインデックスを、昇順で求めるのがこの問題の目的です。

たとえば、入力が heights = [8, 12, 12, 9, 10, 6] の場合、出力は [2, 4, 5] になります。

• インデックス 2(高さ 12):右側の建物は 9・10・6 とすべて低いため、海が見えます。
• インデックス 4(高さ 10):右側には高さ 6 の建物しかなく、海が見えます。
• インデックス 5(高さ 6):右側に建物が存在しないため、必ず海が見えます。

一方、インデックス 1 の高さ 12 の建物は、同じく高さ 12 の建物(インデックス 2)が右側にあるため海は見えません。条件が「より低い」である点、「以下」ではない点に注意しましょう。

解法のアプローチ:単調スタック

この問題は単調スタック(monotonic stack)を使うと効率的に解けます。スタックには「現時点でまだ海が見える可能性のある建物のインデックス」だけを保持し、新しい建物がそれらを遮ったら取り除いていくイメージです。

具体的な手順は次のとおりです。

1. 空のスタックを用意します。
2. 各インデックス idx と高さ h について、次を繰り返します。
   • スタックが空でなく、スタックの末尾のインデックスに対応する建物の高さが h 以下である間、スタックから要素を取り除きます(その建物は新しい建物に遮られ、海が見えなくなるため)。
   • idx をスタックに追加します。
3. 最後にスタックの中身を返します。

実装例(Python)

def solve(heights):
    stack = []
    for idx, h in enumerate(heights):
        while stack and heights[stack[-1]] <= h:
            stack.pop()
        stack.append(idx)
    return stack

heights = [8, 12, 12, 9, 10, 6]
print(solve(heights))

入力と出力

入力: [8, 12, 12, 9, 10, 6]
出力: [2, 4, 5]

動作のトレース

処理がどのように進むか、サンプルデータで確認してみましょう。

最終的にスタックに残っている [2, 4, 5] が答えになります。

計算量

• 時間計算量:O(n) — 各建物は最大でも 1 回 push され、1 回 pop されるだけです。
• 空間計算量:O(n) — 最悪ケース(高さが降順に並んでいる場合)では、すべてのインデックスがスタックに残ります。

各建物ごとに右側をすべて調べる素朴な O(n²) の方法に比べ、単調スタックを使うことで線形時間で解けるのがこの手法の大きなメリットです。