Update trapping-rain-water-ii.cpp

kamyu104 · web-flow · commit 1fa79feddd44 · 2016-09-25T17:57:07.000+08:00
diff --git a/C++/trapping-rain-water-ii.cpp b/C++/trapping-rain-water-ii.cpp
@@ -1,37 +1,45 @@
 // Time:  O(m * n * log(m + n)) ~ O(m * n * log(m * n))
 // Space: O(m * n)
 
-// BFS with priority queue (min heap), refactored version.
 class Solution {
 public:
     /**
      * @param heights: a matrix of integers
      * @return: an integer
      */
     int trapRainWater(vector<vector<int>> &heights) {
-        // Init m_, n_, is_visited_.
         m_ = heights.size();
+        if (!m_) {
+            return 0;
+        }
         n_ = heights[0].size();
+        if (!n_) {
+            return 0;
+        }
+
         is_visited_ = vector<vector<bool>>(m_, vector<bool>(n_, false));
 
         int trap = 0;
 
         // Put the cells on the border into min heap.
         for (int i = 0; i < m_; ++i) {
             heap_.emplace(Cell{i, 0, heights[i][0]});
+            is_visited_[i][0] = true;
             heap_.emplace(Cell{i, n_ - 1, heights[i][n_ - 1]});
+            is_visited_[i][n_ - 1] = true;
         }
         for (int j = 0; j < n_; ++j) {
             heap_.emplace(Cell{0, j, heights[0][j]});
+            is_visited_[0][j] = true;
             heap_.emplace(Cell{m_ - 1, j, heights[m_ - 1][j]});
+            is_visited_[m_ - 1][j] = true;
         }
         const vector<pair<int, int>> directions{{0, -1}, {0, 1},
                                                 {-1, 0}, {1, 0}};
         // BFS with priority queue (min heap)
         while (!heap_.empty()) {
             Cell c = heap_.top();
             heap_.pop();
-            is_visited_[c.i][c.j] = true;
             for (const auto& d : directions) {
             	trap += fill(heights, c.i + d.first, c.j + d.second, c.height);
             }
@@ -49,8 +57,8 @@ class Solution {
 
         // Fill unvisited cell.
         if (!is_visited_[i][j]) {
-            is_visited_[i][j] = true; // Marked as visited.
             heap_.emplace(Cell{i, j, max(height, heights[i][j])});
+            is_visited_[i][j] = true; // Marked as visited.
             return max(0, height - heights[i][j]); // Fill in the gap.
         }
 
