Largest Rectangle in Histogram

from typing import List

class Solution:
    """
    Largest Rectangle in Histogram を解くクラス
    単調スタックによるO(n)時間解法
    """

    def largestRectangleArea(self, heights: List[int]) -> int:
        """
        ヒストグラム内の最大長方形面積を求める

        Args:
            heights: 各棒の高さのリスト
        Returns:
            最大長方形の面積
        Time: O(n), Space: O(n)
        """
        n = len(heights)
        stack = []  # インデックスを格納
        max_area = 0

        # 各要素を処理（番兵として末尾に0を追加）
        for i in range(n + 1):
            # 番兵: i == n の時は高さ0として処理
            current_height = 0 if i == n else heights[i]

            # 現在の高さがスタックトップより低い間、面積計算
            while stack and current_height < heights[stack[-1]]:
                # スタックから高さのインデックスを取得
                height_index = stack.pop()
                height = heights[height_index]

                # 幅を計算: 現在位置 - 左端 - 1
                left_boundary = stack[-1] if stack else -1
                width = i - left_boundary - 1

                # 面積計算と最大値更新
                area = height * width
                max_area = max(max_area, area)

            # 現在のインデックスをスタックにプッシュ
            stack.append(i)

        return max_area

# 使用例
def demo():
    solution = Solution()

    # テストケース1
    heights1 = [2, 1, 5, 6, 2, 3]
    result1 = solution.largestRectangleArea(heights1)
    print(f"Input: {heights1}")
    print(f"Output: {result1}")  # Expected: 10

    # テストケース2
    heights2 = [2, 4]
    result2 = solution.largestRectangleArea(heights2)
    print(f"Input: {heights2}")
    print(f"Output: {result2}")  # Expected: 4

if __name__ == "__main__":
    demo()