LeetCode 84. Largest Rectangle in Histogram

August 19, 2020

We can do this in $O(n)$ time by storing heights on a stack in the format (height, count).

For new larger heights, we just push them to the stack as (height, 1).
For new heights that are smaller than the top of the stack, we remove elements from the stack, amending our max area as we go, and then push our current element (or update the count for the current top of stack if it has the same height).

from typing import List, Tuple
import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def _test(self, i, o):
        self.assertEqual(self.s.largestRectangleArea(i), o)

    def test_examples(self):
        self._test([2,1,5,6,2,3], 10)
        self._test([1,2,5,6,3,1], 10)
        self._test([3,2,4,3,4], 10)

    def test_empty(self):
        self._test([], 0)
        self._test([0], 0)
        self._test([0,0,0,0,0], 0)

    def test_single(self):
        for i in range(1, 11):
            self._test([i], i)
            self._test([0,i,0], i)

    def test_increasing(self):
        self._test([0,1,2], 2)
        self._test([0,1,2,3], 4)
        self._test([0,1,2,3,4], 6)
        self._test([0,1,2,3,4,5], 9)
        self._test([0,1,2,3,4,5,6], 12)

    def test_decreasing(self):
        self._test([2,1,0], 2)
        self._test([3,2,1,0], 4)
        self._test([4,3,2,1,0], 6)
        self._test([5,4,3,2,1,0], 9)
        self._test([6,5,4,3,2,1,0], 12)

class Solution:

    def _processStack(self, stack: List[Tuple[int, int]], min_height: int) -> Tuple[int, int]:
        # Returns (<max area>, <sum of counts>) for the popped elements.
        area = c = 0
        while stack and (min_height == None or stack[-1][0] > min_height):
            b = stack.pop()
            c += b[1]
            area = max(area, b[0] * c)
        return area, c

    def largestRectangleArea(self, heights: List[int]) -> int:
        stack = []
        area = 0
        for i in range(len(heights)):
            poss, c = self._processStack(stack, heights[i])
            area = max(area, poss)
            if stack and stack[-1][0] == heights[i]:
                stack[-1][1] += c+1
            else:
                stack.append([heights[i], c+1])
        poss, _ = self._processStack(stack, None)
        area = max(area, poss)
        return area

if __name__=="__main__":
    unittest.main()