Условие

Дана бинарная матрица mat размером m x n, где 1 представляет солдата, а 0 — гражданского. В каждой строке все 1 расположены перед 0. Строка i считается слабее строки j, если количество солдат в строке i меньше, чем в строке j, или если количество солдат одинаково, но i < j. Необходимо вернуть индексы k самых слабых строк в порядке от самой слабой к самой сильной.

Решение

class Solution {
    fun kWeakestRows(mat: Array<IntArray>, k: Int): IntArray {
        val strength = mat.mapIndexed { index, row ->
            val soldiers = row.count { it == 1 }
            Pair(soldiers, index)
        }
        return strength.sortedWith(compareBy({ it.first }, { it.second }))
            .take(k)
            .map { it.second }
            .toIntArray()
    }
}