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

• 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;
}
};
};
``````

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