🛠️ Union-Find with Sorting || O(NM log(NM)) || ~5ms in Rust & Go
Intuition 🧠To calculate the trapped rainwater, we treat each cell as part of a grid where water flows only through lower-height neighbors. Using a disjoint-set
# Intuition 🧠

To calculate the trapped rainwater, we treat each cell as part of a grid where water flows only through lower-height neighbors. Using a <ins>disjoint-set</ins> (_union-find_) data structure, we group connected components of cells and __track the boundary of the water region__. 
Cells are processed in increasing order of height to ensure we accurately simulate the "filling" of water levels.

# Approach 🚀

1. __Flatten and Sort:__ Represent all cells with their heights and indices, then sort by height.
2. __Union-Find:__ Use a <ins>disjoint-set</ins> to dynamically group connected cells as we process them.
    - __Boundary cells__ are connected to a virtual boundary node.
    - __Neighbors of equal or lower height__ are unioned during processing.
3. __Calculate Volume:__
    - For each height increase, calculate the trapped water by identifying cells not connected to the boundary.
    - Add the water contributed by the height difference for the unconnected "inner" cells.

# Complexity Analysis 🔍

- __Time Complexity:__
    - Sorting cells by height: $O(m × n × \log(m × n))$.
    - <ins>Union-Find</ins> operations (amortized): $O(m × n × α(m × n))$, where $α$ is the <ins>inverse Ackermann function</ins>.
    - __Total:__ $O(m × n × \log(m × n))$ (dominated by sorting).
- __Space Complexity:__
    - <ins>Disjoint-set</ins> data structure: $O(m × n)$.
    - Auxiliary storage (e.g., sorted cells): $O(m × n)$.
    - __Total:__ $O(m × n)$.
    - 
# Code
```rust []
struct DisjointSet {
    parent: Vec<usize>,
    size: Vec<usize>,
}

impl Solution {
    // Directions for moving up, down, left, and right
    const directions: [(i32, i32); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)];

    pub fn trap_rain_water(height_map: Vec<Vec<i32>>) -> i32 {
        let m = height_map.len();
        let n = height_map[0].len();

        // Handle edge cases: if the grid is too small to trap water, return 0
        if m < 3 || n < 3 {
            return 0;
        }

        // Helper function to map a cell's row and column to a unique index in the disjoint set
        let cell_index = |r: usize, c: usize| -> usize { 1 + r * n + c };

        // Helper function to validate a neighbor cell and convert it to its grid coordinates
        let to_valid_cell = |r: i32, c: i32| -> Option<(usize, usize)> {
            if (0..m as i32).contains(&r) && (0..n as i32).contains(&c) {
                Some((r as usize, c as usize))
            } else {
                None
            }
        };

        // Helper function to check if a cell is on the boundary of the grid
        let is_boundary_cell = |r: usize, c: usize| r == 0 || r == m - 1 || c == 0 || c == n - 1;

        // Flatten the grid into a list of cells with their heights and coordinates
        let mut cells = Vec::new();
        for r in 0..m {
            for c in 0..n {
                cells.push((height_map[r][c], r, c));
            }
        }

        // Sort the cells by height in ascending order
        cells.sort_by_key(|&(h, _, _)| h);

        // Initialize the disjoint set with an extra node for the boundary
        let mut ds = DisjointSet::new(m * n + 1);
        let boundary = 0; // Special index representing the boundary node
        let mut volume = 0; // Total trapped water volume
        let mut prev_height = 0; // The previous height processed

        // Process cells in increasing order of height
        for (i, (h, r, c)) in cells.into_iter().enumerate() {
            // Calculate the trapped water when the height increases
            if h > prev_height {
                // Calculate the number of inner cells not connected to the boundary
                let boundary_size = ds.component_size(boundary);
                let inner_cells = (i + 1 - boundary_size) as i32;

                // Add the trapped water for the height difference
                volume += (h - prev_height) * inner_cells;
                prev_height = h;
            }

            // Union the current cell with its valid neighbors of equal or lower height
            for &(dr, dc) in &Self::directions {
                if let Some((nr, nc)) = to_valid_cell(r as i32 + dr, c as i32 + dc) {
                    if height_map[nr][nc] <= h {
                        ds.union(cell_index(r, c), cell_index(nr, nc));
                    }
                }
            }

            // If the current cell is on the boundary, connect it to the boundary node
            if is_boundary_cell(r, c) {
                ds.union(boundary, cell_index(r, c));
            }
        }

        volume
    }
}

impl DisjointSet {
    // Initialize a new disjoint set with `n` elements
    fn new(n: usize) -> Self {
        let parent = (0..n).collect(); // Each element is its own parent
        let size = vec![1; n]; // Each set initially has a size of 1
        Self { parent, size }
    }

    // Find the representative (root) of the set containing `v`
    // and apply path compression for optimization
    fn find(&mut self, v: usize) -> usize {
        if self.parent[v] != v {
            self.parent[v] = self.find(self.parent[v]);
        }
        self.parent[v]
    }

    // Union the sets containing `a` and `b`, attaching the smaller set under the larger set
    fn union(&mut self, a: usize, b: usize) {
        let root_a = self.find(a);
        let root_b = self.find(b);
        if root_a != root_b {
            if self.size[root_a] < self.size[root_b] {
                self.parent[root_a] = root_b;
                self.size[root_b] += self.size[root_a];
            } else {
                self.parent[root_b] = root_a;
                self.size[root_a] += self.size[root_b];
            }
        }
    }

    // Get the size of the set containing `v`
    fn component_size(&mut self, v: usize) -> usize {
        let root = self.find(v);
        self.size[root]
    }
}
```
``` golang []
type DisjointSet struct {
	parent []int
	size   []int
}

func trapRainWater(heightMap [][]int) int {
	directions := [][2]int{{-1, 0}, {1, 0}, {0, -1}, {0, 1}}
	m, n := len(heightMap), len(heightMap[0])

	// Handle edge cases
	if m < 3 || n < 3 {
		return 0 // Trapping water is impossible in grids smaller than 3x3
	}

	// Helper functions
	isValid := func(r, c int) bool {
		return 0 <= r && r < m && 0 <= c && c < n
	}
	isBoundary := func(r, c int) bool {
		return r == 0 || r == m-1 || c == 0 || c == n-1
	}
	cellIndex := func(r, c int) int {
		return 1 + r*n + c // Add 1 to avoid collision with the boundary index
	}

	// Flatten the grid into a list of cells with their row and column indices
	cellsByHeight := make([][2]int, m*n)
	for i := range cellsByHeight {
		cellsByHeight[i][0] = i / n
		cellsByHeight[i][1] = i % n
	}

	// Sort cells by their height in ascending order
	slices.SortFunc(cellsByHeight, func(a, b [2]int) int {
		return heightMap[a[0]][a[1]] - heightMap[b[0]][b[1]]
	})

	// Initialize a disjoint-set structure
	set := NewDisjointSet(m*n + 1)
	boundary := 0 // Special index for the boundary cells

	volume := 0  // Total trapped water volume
	prevH := 0   // Previous processed height

	// Process cells in increasing order of height
	for i := range cellsByHeight {
		r, c := cellsByHeight[i][0], cellsByHeight[i][1]
		h := heightMap[r][c]

		// If the current height is greater than the previous height,
		// calculate the water volume added due to the height difference.
		if h > prevH {
			innerCells := i + 1 - set.size[set.GetId(boundary)] // Inner cells not connected to the boundary
			volume += (h - prevH) * innerCells                  // Add water volume
		}

		// Union boundary cells to the outer space (boundary set)
		if isBoundary(r, c) {
			set.Union(boundary, cellIndex(r, c))
		}

		// Union the current cell with its valid neighbors
		for _, dir := range directions {
			r1, c1 := r+dir[0], c+dir[1]
			if isValid(r1, c1) && heightMap[r1][c1] <= h {
				set.Union(cellIndex(r, c), cellIndex(r1, c1))
			}
		}

		// Update the previous height for the next iteration
		prevH = h
	}

	return volume
}


// Initialize a disjoint-set data structure
func NewDisjointSet(n int) *DisjointSet {
	parent := make([]int, n)
	size := make([]int, n)
	for i := 0; i < n; i++ {
		parent[i] = i
		size[i] = 1
	}
	return &DisjointSet{parent, size}
}

// Find the representative of the set containing element `v`
// and apply path compression.
func (set *DisjointSet) GetId(v int) int {
	for set.parent[v] != v {
		set.parent[v] = set.parent[set.parent[v]] // Path compression
		v = set.parent[v]
	}
	return v
}

// Union two sets containing elements `a` and `b`.
func (set *DisjointSet) Union(a int, b int) int {
	a = set.GetId(a)
	b = set.GetId(b)
	if a == b {
		return a
	}
	// Attach the smaller set under the larger set
	if set.size[a] < set.size[b] {
		a, b = b, a
	}
	set.size[a] += set.size[b]
	set.parent[b] = a
	return a
}
```
['Union Find', 'Go', 'Rust']