파이썬 알고리즘 인터뷰 7장 2번 빗물 트래핑 (리트코드 42번)

Kim Yongbin·2023년 8월 13일

코딩테스트

목록 보기

8/162

Problem

LeetCode - The World's Leading Online Programming Learning Platform

Given n non-negative integers representing an elevation map where the width of each bar is 1, compute how much water it can trap after raining.

각 막대 사이에 얼마나 많은 물이 담길지 계산하여라

Solution

1차 풀이 time out

class Solution:
    @staticmethod
    def strip_by_zero(height: List[int]) -> Optional[List[int]]:
        left = 0
        right = len(height) - 1
        while height[left] == 0 or height[right] == 0:
            if left == right:
                return None
            if height[left] == 0:
                left += 1
            if height[right] == 0:
                right -= 1
        return height[left: right + 1]

    def trap(self, height: List[int]) -> int:
        max_h = max(height)
        water = 0
        for _ in range(max_h):
            height = self.strip_by_zero(height)
            water += height.count(0)
            height = [h-1 if h else h for h in height]
        return water

가로를 x축, 세로를 y축으로 생각했을 때, x축으로 이동하면서 세는 것이 아닌 y축으로 이동하면서 세는 방식으로 접근해보았다.

아래의 과정을 반복하면서 water의 총합을 구하였다.

주어진 리스트의 양의 정수 사이에 존재하는 0 = 해당 층의 water 수
1. 리스트 양쪽 끝의 필요없는 0 slicing
전체 리스트 중 0 이 아닌 정수 -1
반복

문제에 주어진 예제 height = [0,1,0,2,1,0,1,3,2,1,2,1]으로 예시를 들어보겠다.

y = 0, height = [0,1,0,2,1,0,1,3,2,1,2,1]
a. height[2], height[5] ⇒ y=0, water = 2
b. height = [0,0,0,1,0,0,0,2,1,0,1,0]
y = 1, height = [0,0,0,1,0,0,0,2,1,0,1,0]
a. height[4], height[5], height[6], height[10] ⇒ y = 1, water = 4
y = 2, height = [0,0,0,0,0,0,0,1,0,0,0,0]
a. end

따라서, 주어진 height의 총 water = 6이다.

하지만 이 방식은 결국 y축으로 이동하면서 x축으로 스캔하는 방식으로 $O(n^2)$ 의 시간 복잡도를 가지고 있어 Time out이 발생하였다.

2차 풀이 실패

class Solution:
    def count(self, left, height, water):
        right = left + 1
        while right < len(height) and height[left] > height[right]:
            right += 1

        if right >= len(height):
            return left+1, water

        mark_h = min(height[left], height[right])
        if mark_h == 0:
            return right, water

        for h in height[left: right]:
            water += mark_h - h

        return right, water

    def trap(self, height: List[int]) -> int:
        left = 0
        water = 0
        while left < len(height):
            left, water = self.count(left, height, water)
        return water

이번에는 x축으로 이동하는 투 포인터를 이용해보았다. 시작점의 높이와 같거나 높은 높이를 만나면 멈춘 뒤 시작점의 높이, 멈춘 지점의 높이 중 낮은 것을 기준으로 지나온 높이들을 빼서 쌓인 물의 수를 count하였다.

아래의 과정을 반복하면서 water의 총합을 구하였다.

left(시작점)에서부터 height[left] ≤ height[right]가 되는 시점까지 right를 1씩 이동
height[left: right]의 원소들을 min(height[left], height[right])에서 빼서 water의 구하기.
만약 right ≥ len(height)라면 return

이 코드의 경우 height = [4,2,3]과 같이 왼쪽 시작점보다 같거나 높은 height값이 없어 그대로 out되는 케이스를 잡을 수 없었다.

정답 - 투 포인터

from typing import List

class Solution:
    def trap(self, height: List[int]) -> int:
        left, right = 0, len(height) - 1
        left_max, right_max = height[left], height[right]

        water = 0
        while left < right:
            left_max, right_max = max(height[left], left_max), max(height[right], right_max)

            if left_max <= right_max:
                water += left_max - height[left]
                left += 1
            else:
                water += right_max - height[right]
                right -= 1

        return water

왼쪽과 오른쪽에서 각각 접근하면서 현재 높이 최대 값에서 현재의 높이 값을 뺀 값을 각각 더하면서 스캔하는 방식이다. 두 포인터 중 높이가 낮은 곳을 가리키고 있는 포인터가 이동하면서 주어진 heights에서 가장 큰 값을 갖는 원소에서 모이게 된다. 이 방식을 이용하면 시간복잡도 $O(n)$ 으로 풀이가 가능하다.