water trapping 3d problem

water-trapping.cpp

#include <iostream>
#include <vector>
#include <queue>

using namespace std;


struct Node{
    int i,j,h;
    Node(int ii,int jj,int hh):i(ii),j(jj),h(hh){}
    bool operator <(const Node &root) const{
        return h>root.h;
    }
};

int trapRainWater(vector<vector<int>>& heightMap) {
    if(heightMap.size()==0) return 0;
    int m=heightMap.size(),n=heightMap[0].size(),area=0,h=0;
    priority_queue<Node,vector<Node>> q;
    vector<vector<bool>> visit(m,vector<bool>(n,false));
    for(int i=0;i<m;i++){
        for(int j=0;j<n;j++){
            if(i==0||i==m-1||j==0||j==n-1){
                q.push(Node(i,j,heightMap[i][j]));
                visit[i][j]=true;
            }
        }
    }
    vector<int> x={0,0,-1,1},y={-1,1,0,0};
    while(!q.empty()){
        auto f=q.top();
        if(h<f.h) h++;
        else{
            q.pop();
            for(int k=0;k<4;k++){
                int i=f.i+x[k],j=f.j+y[k];
                if(i>=0&&i<m&&j>=0&&j<n&&visit[i][j]==false){
                    int hh=heightMap[i][j];
                    if(hh<h) area+=h-hh;
                    q.push(Node(i,j,hh));
                    visit[i][j]=true;
                }
            }
        }
    }
    return area;
}

int main()
{
    
    vector<vector<int>> map1 = {{123}};
    int result1 = trapRainWater(map1);
    int expected1 = 0;
    cout << (result1 == expected1);
    
    vector<vector<int>> map2 = {{9,2}, {2,1}};
    int result2 = trapRainWater(map2);
    int expected2 = 0;
    cout << (result2 == expected2);

    vector<vector<int>> map3 = {{1,1,1}, {1,1,2}, {1,2,2}};
    int expected3 = 0;
    int result3 = trapRainWater(map3);
    cout << (result3 == expected3);

    vector<vector<int>> map4 = {{6,2,4,123}, {33,2,35,10}, {12,3,0,23}, {83,33,2,34}};
    int expected4 = 2;
    int result4 = trapRainWater(map4);
    cout << (result4 == expected4);

    vector<vector<int>> map5 = {{1,4,3,7,3,2}, {3,2,1,1,3,4}, {5,3,3,5,3,1}};
    int expected5 = 5;
    int result5 = trapRainWater(map5);
    cout << (result5 == expected5);

    vector<vector<int>> map6 = {{13,45,33,77,335,2,3}, {32,25,123,1,3,41,3}, {13,64,19,73,3,212,3}, {52,3,35,35,389,13,3}, {1,4,9,7,3,21,3}};
    int expected6 = 140;
    int result6 = trapRainWater(map6);
    cout << (result6 == expected6);
    
    return 0;
}