C++ priority queue solution


  • 0
    P
    struct Cell {
      int height, i, j;
      Cell(int a, int b, int c) : height(a), i(b), j(c) {}
    };
    
    class Compare {
        public:
        bool operator() (Cell a, Cell b) {return a.height > b.height;}
    };
    
    int trapRainWater(vector<vector<int>>& heightMap) {
        int h = heightMap.size();
        if (h == 0) return 0;
        int w = heightMap[0].size();
        if (w == 0) return 0;
        priority_queue<Cell, vector<Cell>, Compare> pq;
        vector<vector<bool>> visited (h, vector<bool> (w, false));
        for (int i = 0; i < h; i++) {
            pq.push(Cell(heightMap[i][0], i, 0));
            visited[i][0] = true;
            pq.push(Cell(heightMap[i][w - 1], i, w - 1));
            visited[i][w - 1] = true;
        }
        for (int i = 1; i < w - 1; i++) {
            pq.push(Cell(heightMap[0][i], 0, i));
            visited[0][i] = true;
            pq.push(Cell(heightMap[h - 1][i], h - 1, i));
            visited[h - 1][i] = true;
        }
        int ans = 0;
        while (!pq.empty()) {
            Cell cur = pq.top();
            pq.pop();
            if (cur.i > 0 && !visited[cur.i - 1][cur.j]) {
                visited[cur.i - 1][cur.j] = true;
                ans += max(0, cur.height - heightMap[cur.i - 1][cur.j]);
                pq.push(Cell(max(cur.height, heightMap[cur.i - 1][cur.j]), cur.i - 1, cur.j)); 
            }
            if (cur.i < h - 1 && !visited[cur.i + 1][cur.j]) {
                visited[cur.i + 1][cur.j] = true;
                ans += max(0, cur.height - heightMap[cur.i + 1][cur.j]);
                pq.push(Cell(max(cur.height, heightMap[cur.i + 1][cur.j]), cur.i + 1, cur.j)); 
            }
            if (cur.j > 0 && !visited[cur.i][cur.j - 1]) {
                visited[cur.i][cur.j - 1] = true;
                ans += max(0, cur.height - heightMap[cur.i][cur.j - 1]);
                pq.push(Cell(max(cur.height, heightMap[cur.i][cur.j - 1]), cur.i, cur.j - 1)); 
            }
            if (cur.j< w - 1 && !visited[cur.i][cur.j + 1]) {
                visited[cur.i][cur.j + 1] = true;
                ans += max(0, cur.height - heightMap[cur.i][cur.j + 1]);
                pq.push(Cell(max(cur.height, heightMap[cur.i][cur.j + 1]), cur.i, cur.j + 1)); 
            }
        }
        return ans;
    }

Log in to reply
 

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