← Назад

2161.Partition Array According to Given Pivot

Ссылка на задачу: https://leetcode.com/problems/partition-array-according-to-given-pivot/description/


массивы

medium


Описание

Дан 0-indexed массив целых чисел nums и целое число pivot. Необходимо преобразовать nums так, чтобы следующие требования выполнялись:

Верните nums после преобразования.


Примеры

Пример 1

Input:

91251014310

pivot = 10

Output:

95310101214

Элементы 9, 5 и 3 меньше pivot, поэтому они находятся левее 10. Элементы 12 и 14 больше pivot, поэтому они правее 10. Относительный порядок элементов остался неизменным: [9, 5, 3] и [12, 14].

Пример 2

Input:

-3432

pivot = 2

Output:

-3243

Элемент -3 меньше pivot, поэтому он располагается левее. Элементы 4 и 3 больше, поэтому они правее.


Ограничения


Решение

Способ 1

Для реализации нам потребуется три дополнительных массива: before, equal и after. Итерируясь по заданному nums, можно ответить на вопрос, в какой из массивов попадет элемент, сравнив его с pivot. В случае, если он меньше pivot, выбираем массив before, равен - equal, больше - after.

Алгоритм:

  1. Инициализируем три массива: before, equal, after

  2. Сраниваем текущий элемент nums с pivot и кладем в правильный массив

  3. Меняем nums на конкатенацию before, equal и after

Код:

python

class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        before, equal, after = [], [], []
        for i in range(len(nums)):
            if nums[i] < pivot:
                before.append(nums[i])
            elif nums[i] == pivot:
                equal.append(nums[i])
            else:
                after.append(nums[i])
        nums = before + equal + after
        return nums

C++

class Solution {
public:
    vector<int> pivotArray(vector<int>& nums, int pivot) {
        vector<int> before;
        vector<int> equal;
        vector<int> after;

        for (int i = 0; i < nums.size(); ++i) {
            if (nums[i] < pivot) {
                before.push_back(nums[i]);
            } else if (nums[i] == pivot) {
                equal.push_back(nums[i]);
            } else {
                after.push_back(nums[i]);
            }
        }

        vector<int> ans(nums.size());
        int k = 0;
        for (int b : before) {
            ans[k++] = b;
        }
        for (int e : equal) {
            ans[k++] = e;
        }
        for (int a : after) {
            ans[k++] = a;
        }
        nums = ans;

        return nums;
    }
};

Сложность алгоритма по времени - O(N)

Для решения требуется всего один проход по массиву nums размером N и операция конкатенации трех массивов, длина которых в сумме равна N.

Сложность алгоритма по памяти - O(N)

В условии задачи есть некоторое странное требование, чтобы мы возвращали именно nums после всех перестановок. Так как по месту сам массив nums мы не меняем, а для этого используем дополнительный массив (результат конкатенации) размером N, то это можно посчитать как дополнительный объект, необходимый для решения.


Способ 2

Использование трех дополнительных массивов, размер которых растет динамически, в некотором смысле может быть затратно. Их размер нам заранее неизвестен, поэтому какая-то часть времени тратится на реаллокации (да, в C++ можно заранее аллоцировать под каждый из них размер N, так как он является максимальным, но для очень больших N у нас может получиться ситуация, когда довольно большое количество ячеек памяти тратится впустую).

Так, мы можем использовать метод “два указателя”. Если с самого начала инициализировать вспомогательный массив ans размером N, заполненный значениями pivot, то достаточно одной итерации вдоль массива nums двумя указателями (справа и слева) для перезаписи значений, меньших и больших pivot.

Алгоритм:

  1. Инициализировать вспомогательный массив размером N, заполеннный значениями pivot

  2. Инициализировать два указателя: левый (left) с значением 0 и правый (right) с значением N - 1

  3. В рамках одной итерации по массиву перезаписываем ans[left] значениями, меньшими pivot и ans[right] - большими pivot

  4. Присваиваем nums новое значение

Код:

python

class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        ans = [pivot] * len(nums)

        left, right = 0, len(nums) - 1

        i, j = left, right

        while i < len(nums):
            if nums[i] < pivot:
                ans[left] = nums[i]
                left += 1
            if nums[j] > pivot:
                ans[right] = nums[j]
                right -= 1
            i += 1
            j -=1
        nums = ans
        return ans

C++

class Solution {
public:
    vector<int> pivotArray(vector<int>& nums, int pivot) {
        vector<int> ans(nums.size(), pivot);

        int left = 0, right = nums.size() - 1;

        int i = left, j = right;

        while (i < nums.size()) {
            if (nums[i] < pivot) {
                ans[left++] = nums[i];
            }
            if (nums[j] > pivot) {
                ans[right--] = nums[j];
            }
            ++i, --j;
        }
        nums = move(ans);
        return nums;
    }
};

Примечание: для итерации вдоль массива нам тоже требуется два указателя (i и j), так как один указатель бежит слева направо (i), а второй - справа налево (j). Это необходимо для поддержания относительного порядка элементов неизменным.

Сложность алгоритма по времени - O(N)

Время здесь затрачивается всего на один проход по массиву размером N.

Сложность алгоритма по памяти - O(N)

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

Примечание: как таковой оптимизации с точки зрения алгоритма в способе 2 нет, ведь сложность по времени и по памяти ровно та же самая. Тут, скорее, “программистическая” оптимизация.