跳转至

3061. 计算滞留雨水 🔒

题目描述

表:Heights

+-------------+------+
| Column Name | Type |
+-------------+------+
| id          | int  |
| height      | int  |
+-------------+------+
id 是这张表的主键(值互不相同的列),并且保证有序。
这张表的每一行都包含 id 和 height。

编写一个解决方案来计算景观中 沙洲之间 可以滞留的雨水量,认为每个沙洲的 宽度1 个单位。

任何 顺序返回结果表。

结果格式如下例所示。

 

Example 1:

输入: 
Heights table:
+-----+--------+
| id  | height |
+-----+--------+
| 1   | 0      |
| 2   | 1      |
| 3   | 0      |
| 4   | 2      |
| 5   | 1      |
| 6   | 0      |
| 7   | 1      |
| 8   | 3      |
| 9   | 2      |
| 10  | 1      |
| 11  | 2      |
| 12  | 1      |
+-----+--------+
输出: 
+---------------------+
| total_trapped_water | 
+---------------------+
| 6                   | 
+---------------------+
解释: 


上面描绘的高度图(在黑色部分)以图形表示,x 轴表示 id,y 轴表示 heights [0,1,0,2,1,0,1,3,2,1,2,1]。在这个场景中,在蓝色部分滞留了 6 个单位的雨水。

解法

方法一:窗口函数 + 求和

我们使用窗口函数 MAX(height) OVER (ORDER BY id) 来计算每个位置及其左边的最大高度,使用 MAX(height) OVER (ORDER BY id DESC) 来计算每个位置及其右边的最大高度,分别记为 lr。那么每个位置上的蓄水量就是 min(l, r) - height,最后求和即可。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
# Write your MySQL query statement below
WITH
    T AS (
        SELECT
            *,
            MAX(height) OVER (ORDER BY id) AS l,
            MAX(height) OVER (ORDER BY id DESC) AS r
        FROM Heights
    )
SELECT SUM(LEAST(l, r) - height) AS total_trapped_water
FROM T;
1
2
3
4
5
6
7
8
import pandas as pd


def calculate_trapped_rain_water(heights: pd.DataFrame) -> pd.DataFrame:
    heights["l"] = heights["height"].cummax()
    heights["r"] = heights["height"][::-1].cummax()[::-1]
    heights["trapped_water"] = heights[["l", "r"]].min(axis=1) - heights["height"]
    return pd.DataFrame({"total_trapped_water": [heights["trapped_water"].sum()]})

评论