Вернуть отсортированный массив квадратов чисел

Question

На вход дается отсортированный по возрастанию массив. Задача заключается в том, чтобы вернуть отсортированный массив квадратов по возрастанию. На первый взгляд, все очень просто. Пример:

arr = [-5, -3, 0, 1, 3, 6]
result = [0, 1, 9, 9, 25, 36]
# Несколько простых решений (O(N·logN))
result = sorted(el**2 for el in arr)

arr.sort(key=abs)
result = [el**2 for el in arr]

Проблема заключается в том, что нужно сделать это за O(N). Как этого добиться?

Answer 1

я бы делал так (в лоб):

arr = [-5, -3, 0, 1, 3, 6]

left = 0
right = len(arr) - 1

res = []
while left <= right:
    if arr[left]**2 > arr[right]**2:
        res.append(arr[left]**2)
        left += 1
    else:
        res.append(arr[right]**2)
        right -= 1

res.reverse()

print(res)

основной принцип - идем слева направо и справа налево и выбираем самые большие квадраты, в результате у нам получается список квадратов от большего к меньшему

такой подход дает решение задачи за линейное время

Answer 2

Рецепт: разобрать на два списка (отрицательные и остальные), первый список перевернуть, оба возвести в квадрат и слить с помощью heapq.merge.

Работает за линейное время, использует константную дополнительную память:

import heapq

arr = [-5, -3, 0, 1, 3, 6]

result = heapq.merge(
    (v * v for v in reversed(arr) if v <  0), # 9, 25
    (v * v for v in          arr  if v >= 0)  # 0, 1, 9, 36
)

print(*result)

$ python sorted_squares.py
0 1 9 9 25 36

Answer 3

Вот решение без reverse. Оно длиннее, но тут не используются сторонние библиотеки. Воспользуемся свойством отсортированного массива, найдем переход между положительным и отрицательным значением. Далее от него циклами пойдем в разные стороны. Код написан на Java. На собеседовании однажды был такой вопрос и это было правильное решение без использования стороннего кода:

public static int[] sortSquare(int[] sortArr) {
    int l = findChange(sortArr);

    int j = 0;
    int[] resultArr = new int[sortArr.length];

    if (l == -1 && sortArr[0] < 0) {
        for (int i = sortArr.length - 1; i >= 0; i--) {
            resultArr[j] = sortArr[i] * sortArr[i];
            j++;
        }
        return resultArr;
    } else if (l == -1 && sortArr[0] >= 0) {
        for (int k : sortArr) {
            resultArr[j] = k * k;
            j++;
        }
        return resultArr;
    }

    int left = l - 1;
    int right = l;
    while ((left >= 0) && (right <= sortArr.length - 1)) {
        if (Math.abs(sortArr[left]) <= Math.abs(sortArr[right])) {
            resultArr[j] = sortArr[left] * sortArr[left];
            left--;
        } else {
            resultArr[j] = sortArr[right] * sortArr[right];
            right++;
        }
        j++;
    }
    if (right < sortArr.length - 1) {
        for (int i = right; i < sortArr.length; i++) {
            resultArr[j] = sortArr[i] * sortArr[i];
            j++;
        }
    } else {
        for (int i = left; i >= 0; i--) {
            resultArr[j] = sortArr[i] * sortArr[i];
            j++;
        }
    }
    return resultArr;
}

private static int findChange(int[] sortArr) {
    int l = 0;
    int result = -1;

    for (int i = l + 1; i < sortArr.length; i++) {
        if (sortArr[l] < 0 && sortArr[i] >= 0) {
            result = i;
            break;
        } else {
            l = i;
        }
    }
    return result;
}

БЛОГ НА HUSL

Вернуть отсортированный массив квадратов чисел

Ответы (3 шт):