래원

[LeetCode] 407. Trapping Rain Water II (Python) 본문

알고리즘/LeetCode

[LeetCode] 407. Trapping Rain Water II (Python)

Laewon Jeong 2025. 1. 20. 01:37

 

 

Difficulty: Hard

문제 설명


Given an m x n integer matrix heightMap representing the height of each unit cell in a 2D elevation map, return the volume of water it can trap after raining.

 

문제 예제


Example 1:

Input: heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
Output: 4
Explanation: After the rain, water is trapped between the blocks.
We have two small ponds 1 and 3 units trapped.
The total volume of water trapped is 4.

 

Example 2:

Input: heightMap = [[3,3,3,3,3],[3,2,2,2,3],[3,2,1,2,3],[3,2,2,2,3],[3,3,3,3,3]]
Output: 10

 

제한 사항


  • m == heightMap.length
  • n == heightMap[i].length
  • 1 <= m, n <= 200
  • 0 <= heightMap[i][j] <= 2 * 10^4

 

Solution(솔루션)


class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        n = len(heightMap)
        m = len(heightMap[0])
        answer = 0

        visited = defaultdict(bool)
        moves = [[0,1],[0,-1],[1,0],[-1,0]]
        q = []

        for i in range(n):
            for j in range(m):
                if i == 0 or i == n-1 or j == 0 or j == m-1:
                    heapq.heappush(q, (heightMap[i][j], i, j))
                    visited[(i, j)] = True

        while q:
            height, x, y = heapq.heappop(q)

            for move in moves:
                nx = x + move[0]
                ny = y + move[1]

                if 0 <= nx < n and 0 <= ny < m and not visited[(nx, ny)]:
                    visited[(nx, ny)] = True
                    if heightMap[nx][ny] < height:
                        answer += height - heightMap[nx][ny]
                        heightMap[nx][ny] = height
                    heapq.heappush(q, (heightMap[nx][ny], nx, ny))

        return answer

 

처음에는 내부의 작은 height부터 상하좌우의 height을 비교하여 채워질 수 있는 height을 계산했는데, 내부에서 계속 계산을 하려니 답이 보이지 않았다...

 

그래서 가장자리부터 계산을 해봐야겠다고 생각을 했고, 이를 코드로 짜기 시작했다.

 

먼저 q에 가장자리 값을 모두 넣었고, 이를 heapq로 만들었다.

 

이 q를 pop하면 height이 가장 낮은 위치를 얻을 수 있다. 이 위치를 기준으로 상하좌우를 확인하여 인접한 위치가 물이 채워질 수 있는 확인한다.

 

만약 물이 찰 수 있는 위치라면, heightMap[x][y] - heightMap[nx][ny] 값을 answer에 더해주었다.

 

q가 빌 때까지 이 과정을 반복하였고, 마지막에 answer를 return하여 정답을 맞출 수 있었다.

 

생각보다 어려운 문제였다...


문제: 407. Trapping Rain Water II

 

깃허브: github

 

algorithmPractice/LeetCode/407-trapping-rain-water-ii at main · laewonJeong/algorithmPractice

하루 한 문제 챌린지. Contribute to laewonJeong/algorithmPractice development by creating an account on GitHub.

github.com