LeetCode in Kotlin

407. Trapping Rain Water II

Hard

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

Example 1:

Input: heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]

Output: 4

Explanation: 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.

Example 2:

Input: 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]]

Output: 10

Constraints:

Solution

import java.util.PriorityQueue

class Solution {
    private class Cell(val row: Int, val col: Int, val value: Int) : Comparable<Cell?> {
        override operator fun compareTo(other: Cell?): Int {
            return value - other!!.value
        }
    }

    private var water = 0
    private lateinit var visited1: Array<BooleanArray>
    fun trapRainWater(heightMap: Array<IntArray>): Int {
        if (heightMap.size == 0) {
            return 0
        }
        val walls = PriorityQueue<Cell>()
        water = 0
        visited1 = Array(heightMap.size) { BooleanArray(heightMap[0].size) }
        val rows = heightMap.size
        val cols = heightMap[0].size
        // build wall
        for (c in 0 until cols) {
            walls.add(Cell(0, c, heightMap[0][c]))
            walls.add(Cell(rows - 1, c, heightMap[rows - 1][c]))
            visited1[0][c] = true
            visited1[rows - 1][c] = true
        }
        for (r in 1 until rows - 1) {
            walls.add(Cell(r, 0, heightMap[r][0]))
            walls.add(Cell(r, cols - 1, heightMap[r][cols - 1]))
            visited1[r][0] = true
            visited1[r][cols - 1] = true
        }
        // end build wall
        while (walls.isNotEmpty()) {
            val min = walls.poll()
            visit(heightMap, min, walls)
        }
        return water
    }

    private fun visit(height: Array<IntArray>, start: Cell, walls: PriorityQueue<Cell>) {
        fill(height, start.row + 1, start.col, walls, start.value)
        fill(height, start.row - 1, start.col, walls, start.value)
        fill(height, start.row, start.col + 1, walls, start.value)
        fill(height, start.row, start.col - 1, walls, start.value)
    }

    private fun fill(height: Array<IntArray>, row: Int, col: Int, walls: PriorityQueue<Cell>, min: Int) {
        if (row >= 0 && col >= 0 && row < height.size && col < height[0].size &&
            !visited1[row][col]
        ) {
            if (height[row][col] >= min) {
                walls.add(Cell(row, col, height[row][col]))
                visited1[row][col] = true
            } else {
                water += min - height[row][col]
                visited1[row][col] = true
                fill(height, row + 1, col, walls, min)
                fill(height, row - 1, col, walls, min)
                fill(height, row, col + 1, walls, min)
                fill(height, row, col - 1, walls, min)
            }
        }
    }
}