음수가 아닌 정수로 이루어진 길이 n짜리 배열이 있다고 가정해 보겠습니다. 이 배열의 각 값은 높이를 나타내며, 각 막대의 너비는 1입니다. 우리가 구해야 할 것은 비가 온 후 이 지형에 고일 수 있는 물의 총량입니다. 지형은 다음과 같이 표현할 수 있습니다.
위 그림에서 파란색 칸이 8개 있는 것을 확인할 수 있습니다. 따라서 이 경우 출력 결과는 8이 됩니다.
문제 해결 접근 방법
이 문제는 스택(Stack) 자료구조를 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 현재 막대보다 낮은 막대들을 스택에 쌓아두었다가, 더 높은 막대를 만나는 순간 물이 고이는 구간을 계산하는 것입니다.
알고리즘의 동작 과정은 다음과 같습니다.
- 스택 st, 물의 양을 저장할 water := 0, 인덱스 i := 0으로 초기화합니다.
- i가 height 배열의 길이보다 작은 동안 다음을 반복합니다.
- 스택이 비어 있거나, 스택 최상단 인덱스의 높이가 현재 높이(height[i])보다 크거나 같으면 현재 인덱스 i를 스택에 push하고 i를 1 증가시킵니다.
- 그렇지 않은 경우:
- x := 스택 최상단 요소를 꺼내고(pop), 해당 요소를 제거합니다.
- 스택이 비어 있지 않다면:
- temp := height[스택 최상단]과 height[i] 중 작은 값
- dist := i − 스택 최상단 인덱스 − 1
- water := water + dist × (temp − height[x])
- 반복이 끝나면 water 값을 반환합니다.
파이썬 구현 코드
아래 코드를 통해 실제 구현 방법을 더 잘 이해할 수 있습니다.
class Solution(object): def trap(self, height): stack = [] water = 0 i=0 while i<len(height): if len(stack) == 0 or height[stack[-1]]>=height[i]: stack.append(i) i+=1 else: x = stack[-1] stack.pop() if len(stack) != 0: temp = min(height[stack[-1]],height[i]) dist = i - stack[-1]-1 water += dist*(temp - height[x]) return water ob = Solution() print(ob.trap([2,5,2,0,5,8,8]))
입력
[2,5,2,0,5,8,8]
출력
8
동작 원리 정리
이 알고리즘에서 각 인덱스는 최대 한 번 push되고 한 번 pop되므로 시간 복잡도는 O(n)입니다. 공간 복잡도 역시 최악의 경우 모든 인덱스가 스택에 저장될 수 있으므로 O(n)입니다. 두 포인터(Two Pointer) 기법을 사용하면 공간 복잡도를 O(1)까지 줄일 수 있지만, 스택 기반 풀이는 직관적이라 코딩 테스트에서 널리 활용됩니다.