[LeetCode]42. Trapping Rain Water
题目描述
思路
维护一个栈
先将数组中第一个元素入栈,
后续元素和栈顶元素做比较,
若小于栈顶元素,则入栈,继续循环
若大于栈顶元素,则栈顶元素出栈,计算增加的数目,保留当前栈顶元素值
重复与栈顶元素比较,直到满足小于条件或栈中元素为空
举例来说
对于 heights = {2,1,0,1,3}
开始将2入栈,此时栈中元素为{2}
1 < 2,所以1入栈,此时栈中元素为{2, 1}
0 < 1,所以0入栈,此时栈中元素为{2, 1, 0}
1 > 0,此时计算要增加的数目
- 初始化added = 0, count = 0,
- 计算要增加的数目 added += heights[i] - s.top() = 1
- 计数count++,此时count = 1
- 将0出栈,栈中元素为{2, 1}
- head = s.top() = 1
此时子循环结束 计算相加 - 若head >= heights[i],则表明计算的added可以直接加入到res中 、
- 若head < heights[i],则表明计算的added多算了 count * (heights[i] -
head),因此,added需要减去多算的部分才可以加入res
相加计算结束,继续判断 - 若head >= heights[i],则表明仍然需要计算结果,需要把刚刚出栈的数再入栈,防止重复计算结果,将入栈的数同一按照 heights[i] 大小入栈
入栈判断结束
最后将heights[i]入栈
后续循环过程省略,详见代码
代码
#include <iostream>
#include <vector>
#include <stack>
using namespace std;
class Solution {
public:
int trap(vector<int>& heights) {
if (heights.size() == 0)
return 0;
stack<int> s;
s.push(heights[0]);
int count = 0, added = 0, res = 0, head = 0;
for (int i = 1; i < heights.size(); ++i) {
count = 0, added = 0, head = 0;
if (heights[i] > s.top()) {
while (s.size() && heights[i] >= s.top()) {
head = s.top();
if (s.size() != 1) {
added += heights[i] - head;
++count;
}
s.pop();
}
if (s.size())
head = s.top();
if (s.size() || head) {
added = head >= heights[i] ? added : added - count * (heights[i] - head);
res += added;
if (head >= heights[i]) {
while (count--) {
s.push(heights[i]);
}
}
}
}
s.push(heights[i]);
}
return res;
}
};
int main() {
vector<int> heights = { 2,2,1,0,1,3 };
Solution s;
cout << s.trap(heights) << endl;
system("pause");
return 0;
}