Algorithms : Trapping Rain Water II

The problem is from leetcode 407.

https://leetcode-cn.com/problems/trapping-rain-water-ii

Given an m x n matrix of positive integers representing the height of each unit cell in a 2D elevation map, compute the volume of water it is able to trap after raining.

 

Note:

Both m and n are less than 110. The height of each unit cell is greater than 0 and is less than 20,000.

 

Example:

Given the following 3x6 height map:
[
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]

Return 4.

The above image represents the elevation map [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]] before the rain.

After the rain, water is trapped between the blocks. The total volume of water trapped is 4.

The first solution

Using the DFS. 

class Solution {
private:
    struct cell {
        int row;
        int col;
        int height;
    };
    static constexpr int MAX_INTEGER = 20001;
    // search for wall and record the path.
    int DFS(vector<vector<int>> & heightMap, int i, int j, int height,
            vector<cell> & trackPaths) {
        int rows = heightMap.size();
        int cols = heightMap[0].size();
        if(i < 0 || i >= rows || j < 0 || j >= cols) return -1; // we hit the edge.
        else if(heightMap[i][j] < height) return -1; // there is a lower stage.
        else if(heightMap[i][j] > height) return heightMap[i][j]; // hit a wall , just return the height.
        else { // the same level
            trackPaths.push_back(cell{i, j, heightMap[i][j]}); // record the current path.
            // set as visited.
            heightMap[i][j] = MAX_INTEGER;
            int a1, a2, a3, a4;
            a1 = DFS(heightMap, i + 1, j, height, trackPaths);
            if(a1 == -1) return -1;
            a2 = DFS(heightMap, i - 1, j, height, trackPaths);
            if(a2 == -1) return -1;
            a3 = DFS(heightMap, i, j + 1, height, trackPaths);
            if(a3 == -1) return -1;
            a4 = DFS(heightMap, i, j - 1, height, trackPaths);
            if(a4 == -1) return -1;
            int h = min(min(a1, a2), min(a3, a4));
            return h;
        }
    }
    
    void recoverPath(vector<vector<int>>& heightMap, vector<cell> & trackPaths) {
        for(auto &entry : trackPaths) {
            heightMap[entry.row][entry.col] = entry.height;
        }
    }
    
    int setPath(vector<vector<int>>& heightMap, vector<cell> & trackPaths, int height) {
        int vol = 0;
        for(auto &entry : trackPaths) {
            heightMap[entry.row][entry.col] = height;
            vol += height - entry.height; // accumulate.
        }
        return vol;     
    }
public:
    int trapRainWater(vector<vector<int>>& heightMap) {
        int vol = 0;
        int rows = heightMap.size();
        if(rows == 0) return 0;
        int cols = heightMap[0].size();
        for(int i = 0; i < rows; ++i) {
            for (int j = 0; j < cols; ++j) {
                // try to fill with water.
                while (true) {
                    vector<cell> path;
                    int currentHeight = heightMap[i][j];
                    int targetHeight = DFS(heightMap, i, j, currentHeight, path);
                    if(targetHeight == -1) {
                        recoverPath(heightMap, path);
                        break;
                    } else {
                        vol += setPath(heightMap, path, targetHeight);
                    }
                }   
            }
        }
        return vol;
    }
};

The running time is  148ms with 43.2MB memory

 

The optimized solution : using the min heap with BFS

class Solution {
private:
    struct cell {
        int row;
        int col;
        int height;
    };
    struct cmp {
        bool operator()(const cell & l, const cell & r){
            return l.height > r.height;
        }
    };
public:
    int trapRainWater(vector<vector<int>>& heightMap) {
        if (heightMap.size() == 0 || heightMap.size() < 3 || heightMap[0].size() < 3)
            return 0;
        int rows = heightMap.size();
        int cols = heightMap[0].size();
        
        // minHeap to scan.
        priority_queue<cell, vector<cell>, cmp> minHeap;
        // indicate that whether a cell is visited.
        vector<vector<bool>> visited(rows, vector<bool>(cols, false)); 
        // add the left/right boundary to the minHeap.
        for(int i = 0; i < rows; ++i) {
            visited[i][0] = visited[i][cols - 1] = true;
            minHeap.push(cell{i, 0, heightMap[i][0]});
            minHeap.push(cell{i, cols - 1, heightMap[i][cols - 1]});
        }     
        // add the up/down boundary to the minHeap.
        for(int j = 0; j < cols; ++j) {
            visited[0][j] = visited[rows - 1][j] = true;
            minHeap.push(cell{0, j, heightMap[0][j]});
            minHeap.push(cell{rows - 1, j, heightMap[rows - 1][j]});
        }
        
        // direction map
        constexpr int direction[][2] = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
        int vol = 0;
        while(!minHeap.empty()) {
            auto Cell = minHeap.top();
            for(int i = 0; i < 4; ++i) {
                int m = Cell.row + direction[i][0];
                int n = Cell.col + direction[i][1];
                // the target cell is legal and is has not been visited.
                if(m >= 0 && m < rows && n >=0 && n < cols && !visited[m][n]) {
                    vol += max(0, Cell.height - heightMap[m][n]);
                    visited[m][n] = true;
                    minHeap.push(cell{m, n, max(Cell.height, heightMap[m][n])});
                }
            }
            minHeap.pop();
        }
        return vol;
    }
};

The running time is 52ms with 10.8MB

The total code with the test case is :

#include <algorithm>
#include <iostream>
#include <list>
#include <queue>
#include <string>
#include <vector>

using namespace std;

class Solution
{
private:
    struct cell
    {
        int row;
        int col;
        int height;
    };
    struct cmp
    {
        bool operator()(const cell &l, const cell &r)
        {
            return l.height > r.height;
        }
    };

public:
    int trapRainWater(vector<vector<int>> &heightMap)
    {
        if (heightMap.size() == 0 || heightMap.size() < 3 || heightMap[0].size() < 3)
            return 0;
        int rows = heightMap.size();
        int cols = heightMap[0].size();

        // minHeap to scan.
        priority_queue<cell, vector<cell>, cmp> minHeap;
        // indicate that whether a cell is visited.
        vector<vector<bool>> visited(rows, vector<bool>(cols, false));
        // add the left/right boundary to the minHeap.
        for (int i = 0; i < rows; ++i)
        {
            visited[i][0] = visited[i][cols - 1] = true;
            minHeap.push(cell{i, 0, heightMap[i][0]});
            minHeap.push(cell{i, cols - 1, heightMap[i][cols - 1]});
        }
        // add the up/down boundary to the minHeap.
        for (int j = 0; j < cols; ++j)
        {
            visited[0][j] = visited[rows - 1][j] = true;
            minHeap.push(cell{0, j, heightMap[0][j]});
            minHeap.push(cell{rows - 1, j, heightMap[rows - 1][j]});
        }

        // direction map
        constexpr int direction[][2] = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
        int vol = 0;
        while (!minHeap.empty())
        {
            auto Cell = minHeap.top();
            for (int i = 0; i < 4; ++i)
            {
                int m = Cell.row + direction[i][0];
                int n = Cell.col + direction[i][1];
                if (m == 2 && n == 1) {
                    int tmp = 0;
                    ++tmp;
                }
                // the target cell is legal and is has not been visited.
                if (m >= 0 && m < rows && n >= 0 && n < cols && !visited[m][n])
                {
                    vol += max(0, Cell.height - heightMap[m][n]);
                    visited[m][n] = true;
                    minHeap.push(cell{m, n, max(Cell.height, heightMap[m][n])});
                }
            }
            minHeap.pop();
        }
        return vol;
    }
};

int main()
{
    /*vector<vector<int>> heightMap = {
        vector<int>{9, 9, 9, 9, 9, 9, 8, 9, 9, 9, 9}, 
        vector<int>{9, 0, 0, 0, 0, 0, 1, 0, 0, 0, 9}, 
        vector<int>{9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}, 
        vector<int>{9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}, 
        vector<int>{9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9}
    };*/

    vector<vector<int>> heightMap = {
        vector<int>{5, 8, 7, 7}, 
        vector<int>{5, 2, 1, 5}, 
        vector<int>{7, 1, 7, 1}, 
        vector<int>{8, 9, 6, 9}, 
        vector<int>{9, 8, 9, 9}
    };

    Solution sol;
    int vol = sol.trapRainWater(heightMap);
    cout << vol << endl;
    return 0;
}