Oтсортированные квадраты

Постановка задачи

nums = [-4, -1, 0, 3, 10]

[0, 1, 9, 16, 100]

nums = [-7, -3, 2, 3, 11]

[4, 9, 9, 49, 121]

Решение

Мы продолжаем погружение в метод двух указателей. Эта задача считается классикой уровня Easy на LeetCode. Она идеально подходит для того, чтобы закрепить базовый навык: умение эффективно обрабатывать массив, используя два указателя.

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

Самый очевидный способ: возвести все числа в квадрат, а затем отсортировать массив. Сложность такого решения O(n log n) из-за сортировки. Здесь мы никак не используем то, что исходный массив уже отсортирован, и интервьюер ожидает от вас ответа на дополнительный вопрос: можете ли вы найти решение со сложностью O(n).

Идея решения:

Попробуйте пару минут самостоятельно подумать, как подход с двумя указателями может помочь решить эту задачу за O(n). Если найти решение оказалось просто - отлично. Если возникли сложности, сейчас мы разберем идею.

Начнем с того, что в этой задаче от нас не требуется решение "на месте" (in-place). Мы можем позволить себе выделить дополнительную память под новый массив той же длины, что и исходный.

Далее, если бы массив не содержал отрицательных чисел, задача решалась бы простым проходом по массиву и возведением в квадрат каждого элемента. Наличие отрицательных чисел все меняет.

Теперь ответим на вопрос: где в исходном массиве могут находиться самые большие значения квадратов?

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

Если посмотреть на края исходного массива, можно заметить: максимальное значение квадрата всегда находится либо в самом конце, где стоит самое большое положительное число, либо в самом начале, где стоит самое большое по модулю отрицательное число.

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

Основные шаги:

Вот основные шаги алгоритма:

1. Инициализация:

   • Создаем новый массив res того же размера, что и исходный.
   • Ставим индекс j на последнюю позицию в res, так как мы будем заполнять его от самых больших значений к меньшим.
   • Устанавливаем два указателя: левый l на начало массива nums и правый r на его конец.

2. В цикле: пока левый указатель не пересечет правый l <= r, сравниваем числа по краям.

   • Для сравнения берем их значения по модулю, так как минус при возведении в квадрат исчезнет.
   • Если число слева по модулю больше - возводим его в квадрат, записываем в res[j] и сдвигаем левый указатель вправо (l++).
   • Иначе, если правое число больше или они равны, возводим в квадрат правое число, записываем в res[j] и сдвигаем правый указатель влево (r--).
   • После каждой записи уменьшаем индекс j на единицу, подготавливая место для следующего квадрата.

3. Завершение цикла:

   • Когда указатели встретились и обработали последний элемент, цикл завершается, и мы возвращаем полностью заполненный массив res.

Разберем пример:

nums = [-4, -1, 0, 3, 10]

0. Инициализация

   • l = 0, r = 4, j = 4
   • res = [0, 0, 0, 0, 0]

1. Итерация 1

   • Сравниваем: abs(nums[0]) и abs(nums[4]) -> abs(-4) и abs(10).
   • 10 > 4, значит берем правый элемент.
   • res[4] = 10 * 10 = 100
   • Сдвигаем r (стал 3) и j (стал 3).
   • Массив: res = [0, 0, 0, 0, 100]

2. Итерация 2

   • Сравниваем: abs(nums[0]) и abs(nums[3]) -> abs(-4) и abs(3).
   • 4 > 3, значит берем левый элемент.
   • res[3] = -4 * -4 = 16
   • Сдвигаем l (стал 1) и j (стал 2).
   • Массив: res = [0, 0, 0, 16, 100]

3. Итерация 3

   • Сравниваем: abs(nums[1]) и abs(nums[3]) -> abs(-1) и abs(3).
   • 3 > 1, значит берем правый элемент.
   • res[2] = 3 * 3 = 9
   • Сдвигаем r (стал 2) и j (стал 1).
   • Массив: res = [0, 0, 9, 16, 100]

4. Итерация 4

   • Сравниваем: abs(nums[1]) и abs(nums[2]) -> abs(-1) и abs(0).
   • 1 > 0, значит берем левый элемент.
   • res[1] = -1 * -1 = 1
   • Сдвигаем l (стал 2) и j (стал 0).
   • Массив: res = [0, 1, 9, 16, 100]

5. Итерация 5 (финальная, l == r)

   • Сравниваем: abs(nums[2]) и abs(nums[2]) -> abs(0) и abs(0).
   • Значения равны, берем правый элемент.
   • res[0] = 0 * 0 = 0
   • Сдвигаем r (стал 1) и j (стал -1).
   • Массив: res = [0, 1, 9, 16, 100]

Итог: [0, 1, 9, 16, 100]

Оценка сложности

Временная сложность

Мы за один проход по массиву nums обработали все числа, возвели их в квадрат и расположили в правильных местах в результате. Поэтому временная сложность равна O(n).

Пространственная сложность

Мы создаем массив res для результата, поэтому пространственная сложность равна O(n).

Код решения

Приведем код решения.

func sortedSquares(nums []int) []int {
    l := 0
    r := len(nums) - 1
    j := len(nums) - 1
    res := make([]int, len(nums))
    for l <= r {
        s := nums[r]
        if abs(nums[l]) > abs(s) {
            s = nums[l]
            l++
        } else {
            r--
        }
        res[j] = s * s
        j--
    }
    return res
}

Итоги

1. Главная идея. Мы использовали тот факт, что исходный массив отсортирован. Это позволило нам применить метод двух указателей и сократить сложность до линейной.

2. Направление движения. Поскольку самые большие квадраты находятся по краям (в начале и в конце массива), мы заполняли результирующий массив с конца к началу. Это позволило нам сразу ставить максимальные значения на их финальные позиции.

3. Новый массив. Если в условии явно не сказано, что нужно решать задачу на месте, мы можем создать новый массив для результата. Это делает решение проще и не меняет входные данные.

4. Критический момент: условие l <= r.

   • Использование знака "меньше или равно" принципиально важно.
   • Когда указатели l и r встречаются на одном и том же элементе, это означает, что в исходном массиве остался последний (самый маленький по модулю) необработанный элемент.
   • Если написать просто l < r, этот последний элемент не попадет в цикл, и мы вернем массив с пустым (нулевым) первым элементом.