Ссылка на задачу: https://leetcode.com/problems/partition-array-according-to-given-pivot/description/
Дан 0-indexed массив целых чисел nums и целое число pivot. Необходимо преобразовать nums так, чтобы следующие требования выполнялись:
каждый элемент, меньший pivot, находился левее чем те элементы, что больше pivot
каждый элемент, равный pivot, находился между элементами, меньшими и большими pivot
относительный порядок элементов, меньших pivot и больших pivot, сохранялся неизменным
Верните nums после преобразования.
Input:
9 | 12 | 5 | 10 | 14 | 3 | 10 |
---|
pivot = 10
Output:
9 | 5 | 3 | 10 | 10 | 12 | 14 |
---|
Элементы 9, 5 и 3 меньше pivot, поэтому они находятся левее 10. Элементы 12 и 14 больше pivot, поэтому они правее 10. Относительный порядок элементов остался неизменным: [9, 5, 3] и [12, 14].
Input:
-3 | 4 | 3 | 2 |
---|
pivot = 2
Output:
-3 | 2 | 4 | 3 |
---|
Элемент -3 меньше pivot, поэтому он располагается левее. Элементы 4 и 3 больше, поэтому они правее.
1 <= nums.length <= 10^5
-10^6 <= nums[i] <= 10^6
pivot равен одному из элементов nums
Для реализации нам потребуется три дополнительных массива: before, equal и after. Итерируясь по заданному nums, можно ответить на вопрос, в какой из массивов попадет элемент, сравнив его с pivot. В случае, если он меньше pivot, выбираем массив before, равен - equal, больше - after.
Алгоритм:
Инициализируем три массива: before, equal, after
Сраниваем текущий элемент nums с pivot и кладем в правильный массив
Меняем 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, то это можно посчитать как дополнительный объект, необходимый для решения.
Использование трех дополнительных массивов, размер которых растет динамически, в некотором смысле может быть затратно. Их размер нам заранее неизвестен, поэтому какая-то часть времени тратится на реаллокации (да, в C++ можно заранее аллоцировать под каждый из них размер N, так как он является максимальным, но для очень больших N у нас может получиться ситуация, когда довольно большое количество ячеек памяти тратится впустую).
Так, мы можем использовать метод “два указателя”. Если с самого начала инициализировать вспомогательный массив ans размером N, заполненный значениями pivot, то достаточно одной итерации вдоль массива nums двумя указателями (справа и слева) для перезаписи значений, меньших и больших pivot.
Алгоритм:
Инициализировать вспомогательный массив размером N, заполеннный значениями pivot
Инициализировать два указателя: левый (left) с значением 0 и правый (right) с значением N - 1
В рамках одной итерации по массиву перезаписываем ans[left] значениями, меньшими pivot и ans[right] - большими pivot
Присваиваем 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 нет, ведь сложность по времени и по памяти ровно та же самая. Тут, скорее, “программистическая” оптимизация.