Квадраты упорядоченного массива

Условие

Дан массив целых чисел arr длиной length, отсортированный по возрастанию. Требуется реализовать функцию, которая формирует новый массив, содержащий квадраты всех элементов исходного массива, и при этом сохраняет порядок возрастания в результирующей последовательности.

Пример: [-3, 2, 4] → [4, 9, 16]

Ответ

Сравнивайте квадраты крайних элементов и записывайте больший с конца результирующего массива. Это даёт O(n) времени без дополнительной сортировки.

fun sortedSquares(arr: IntArray): LongArray {
    val result = LongArray(arr.size)
    var left = 0
    var right = arr.lastIndex
    for (write in result.lastIndex downTo 0) {
        val a = arr[left].toLong() * arr[left]
        val b = arr[right].toLong() * arr[right]
        if (a > b) {
            result[write] = a
            left++
        } else {
            result[write] = b
            right--
        }
    }
    return result
}

Максимальный модуль среди оставшихся чисел находится на одном из концов. После удаления этого элемента правило сохраняется. Использование Long до умножения защищает от переполнения для любого входного Int. Для пустого массива цикл не выполняется. Память O(n) под результат.