(C++) The priority queue sol works for both trap water I & II


  • 0
    M

    My following 2D (C++) solution using priority queue is due to yuhaowang001's post.

    I found the same idea also applies to the 1D problem - Still use priority queue but with only two directions (left or right).

    2D sol:

    class Solution {
    public:
        int trapRainWater(vector<vector<int>>& heightMap) {
            if(heightMap.size()==0 || heightMap[0].size()==0) return 0;
            
            //enque boundary cells & mark them as visited
            int m=heightMap.size(), n=heightMap[0].size();
            priority_queue<Cell> que;
            vector<vector<bool>> visited(m, vector<bool>(n, false));
            for(int i=0; i<m; i++) {
                que.push(Cell(i,0,heightMap[i][0]));
                que.push(Cell(i,n-1,heightMap[i][n-1]));
                visited[i][0]=true;
                visited[i][n-1]=true;
            }
            for(int j=0; j<n; j++) {
                que.push(Cell(0,j,heightMap[0][j]));
                que.push(Cell(m-1,j,heightMap[m-1][j]));
                visited[0][j]=true;
                visited[m-1][j]=true;
            } 
            
            // Going inward - using min_height node in queue & move in 4 directions & trap water locally
            vector<int> xs({0,0,1,-1}), ys({1,-1,0,0}); // 4 directions
            int water = 0;
            while(!que.empty()) {
                Cell cell = que.top();
                que.pop();
                for(int k=0; k<4; k++) {//move in 4 directions, if possible
                    int xx=cell.x+xs[k], yy=cell.y+ys[k];
                    if(xx<0 || xx>=m || yy<0 || yy>=n || visited[xx][yy]) continue;
                    water+=max(cell.h-heightMap[xx][yy],0);
                    que.push(Cell(xx,yy,max(cell.h,heightMap[xx][yy])));    
                    visited[xx][yy]=true;
                }
            }
            return water;
            
        }
    private:
        struct Cell {
            int x, y, h;
            Cell(int xx, int yy, int hh):x(xx), y(yy), h(hh) {}
            bool operator<(const Cell& rhs) const {
                return this->h > rhs.h;
            }
        };
    };
    

    1D sol:

    class Solution {
    public:
        int trap(vector<int>& height) {
            const int n = height.size();
            if(n==0) return 0;
            
            priority_queue<Cell> que;
            vector<bool> visited(n, false);
            //enque leftmost & rightmost cell
            que.push(Cell(0,height[0]));
            que.push(Cell(n-1,height[n-1]));
            visited[0]=true;
            visited[n-1]=true;
            
            int water = 0;
            vector<int> xs({1,-1}); // right, left
            while(!que.empty()) {
                Cell cell = que.top(); 
                que.pop();
                for(int inc:xs) {
                    int xx = cell.x+inc;
                    if(xx<0 || xx>=n || visited[xx]) continue;
                    water+=max(0,cell.h-height[xx]);
                    que.push(Cell(xx,max(cell.h,height[xx])));
                    visited[xx]=true;
                }
            }
            return water;    
            
        }
    private:
        struct Cell {
            int x, h;
            Cell(int xx, int hh):x(xx), h(hh) {}
            bool operator<(const Cell& rhs) const {
                return this->h > rhs.h;
            }
        };
    };
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.