Matrix Cells in Distance Order

To solve this easy LC problem, my initial idea was to use a TreeMap<Int, Pair<Int, Int>> to store the distance and its corresponding coordinate. TreeMap automatically sorts the keys during insertion so that will save us some time instead of sorting the keys ourselves.

But when I run the code, it failed in some test cases. I realized that of course, some distances will be the same and TreeMap does not allow duplicate keys. Therefore, only one of the coordinates with equal distances to the target coordinate will be inserted in the map.

The fix is easy. Simply change the type parameter of TreeMap value to an ArrayList so that we can chain all these coordinates with the same value within the same bucket (distance).

import java.util.*

class Solution {
    fun allCellsDistOrder(R: Int, C: Int, r0: Int, c0: Int): Array<IntArray> {
        val distance: (Pair<Int, Int>, Pair<Int, Int>) -> Int = { p1, p2 ->
            Math.abs(p1.first - p2.first) + Math.abs(p1.second - p2.second)
        }
        
        var output = TreeMap<Int, ArrayList<Pair<Int, Int>>>()
        val target = Pair(r0, c0)
        
        (0 until R).forEach { row ->
            (0 until C).forEach { column ->
                val coordPair = Pair(row, column)
                val dist = distance(coordPair, target)
                if(!output.containsKey(dist)){
                    output.put(dist, ArrayList<Pair<Int, Int>>())
                }
                
                output.get(dist)?.add(coordPair)
            }            
        }
        
        return output.flatMap {it.value.map { intArrayOf(it.first, it.second) } }.toTypedArray()
    }
}

Runtime: 376 ms, faster than 70.00% of Kotlin online submissions for Matrix Cells in Distance Order.
Memory Usage: 39.8 MB, less than 100.00% of Kotlin online submissions for Matrix Cells in Distance Order.

Marc Lester Tan