[LeetCode]42. Trapping Rain Water

[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;
}