<p>Given an <code>m x n</code> integer matrix <code>heightMap</code> representing the height of each unit cell in a 2D elevation map, return <em>the volume of water it can trap after raining</em>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2021/04/08/trap1-3d.jpg" style="width: 361px; height: 321px;" />
<pre>
<strong>Input:</strong> heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
<strong>Output:</strong> 4
<strong>Explanation:</strong> After the rain, water is trapped between the blocks.
We have two small ponds 1 and 3 units trapped.
The total volume of water trapped is 4.
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2021/04/08/trap2-3d.jpg" style="width: 401px; height: 321px;" />
<pre>
<strong>Input:</strong> heightMap = [[3,3,3,3,3],[3,2,2,2,3],[3,2,1,2,3],[3,2,2,2,3],[3,3,3,3,3]]
<strong>Output:</strong> 10
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>m == heightMap.length</code></li>
	<li><code>n == heightMap[i].length</code></li>
	<li><code>1 &lt;= m, n &lt;= 200</code></li>
	<li><code>0 &lt;= heightMap[i][j] &lt;= 2 * 10<sup>4</sup></code></li>
</ul>

class Solution {
public:
    int trapRainWater(vector<vector<int>>& heightMap) {
        int m = heightMap.size();
        int n = heightMap[0].size();
        int total_cells = m * n;
        vector<int> heights(total_cells);
        
        for (int i = 0; i < m; ++i) {
            const vector<int>& row = heightMap[i];
            for (int j = 0; j < n; ++j) {
                heights[i*n + j] = row[j];
            }
        }
        std::sort(heights.begin(), heights.end());
        auto last = std::unique(heights.begin(), heights.end());

        int result = 0;
        int last_height = 0;
        bool* wall_grid = new bool[total_cells];
        bool* visited = new bool[total_cells] {false};
        for (auto iter = heights.begin(); iter != last; ++iter) {
            int current_height = *iter;
            if (current_height <= last_height) continue;
            int level_height = current_height - last_height; 

            // update grid
            for (int i = 0; i < m; ++i) {
                const vector<int>& row = heightMap[i];
                for (int j = 0; j < n; ++j) {
                    // false is empty, true is wall
                    wall_grid[i*n + j] = row[j] > last_height;
                }
            }


            // traverse the grid
            for (int i = 0; i < m; ++i) {
                for (int j = 0; j < n; ++j) {
                    if (!wall_grid[i*n + j] && !visited[i*n + j]) {
                        bool reached_edge = false;
                        int cells = 0;
                        traverse(i, j, m, n, wall_grid, visited, reached_edge, cells);
                        if (!reached_edge) result += cells * level_height;
                    }
                }
            }

            last_height = current_height;
            // clear visited
            memset(visited, false, total_cells);
        }

        delete [] wall_grid;
        delete [] visited;

        return result;
    }

    void traverse(int row, int col, int m, int n, bool* grid, bool* visited, bool& reached_edge, int& cells) {
        cells += 1;
        if (row == 0 || row == m - 1 || col == 0 || col == n - 1) reached_edge = true;
        visited[row*n + col] = true;
        if (row > 0 && !grid[(row - 1)*n + col] && !visited[(row - 1)*n + col]) {
            traverse(row - 1, col, m, n, grid, visited, reached_edge, cells);
        }
        if (row < m-1 && !grid[(row + 1)*n + col] && !visited[(row + 1)*n + col]) {
            traverse(row + 1, col, m, n, grid, visited, reached_edge, cells);
        }
        if (col > 0 && !grid[row*n + col - 1] && !visited[row*n + col - 1]) {
            traverse(row, col - 1, m, n, grid, visited, reached_edge, cells);
        }
        if (col < n-1 && !grid[row*n + col + 1] && !visited[row*n + col + 1]) {
            traverse(row, col + 1, m, n, grid, visited, reached_edge, cells);
        }
    }
};
23256