42. 接雨水

exiaohu 于 2021-11-03 发布

题目链接:42. 接雨水

假定给定范围内的雨水已经接满:

使用优先队列,从最低的位置开始蓄水,知道蓄满给定范围。

import heapq
from typing import List


class Solution:
    def trap(self, height: List[int]) -> int:
        n = len(height)
        if n <= 2:
            return 0

        q, visited, ans = [], {0, n - 1}, 0
        heapq.heappush(q, (height[0], 0))
        heapq.heappush(q, (height[n - 1], n - 1))

        while len(q) > 0:
            h, x = heapq.heappop(q)
            for dx in (-1, 1):
                i = x + dx
                if 0 <= i < n and i not in visited:
                    visited.add(i)
                    wh = max(h, height[i])
                    ans += wh - height[i]
                    heapq.heappush(q, (wh, i))

        return ans