Сортировка цветов

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

nums = [2, 0, 2, 1, 1, 0]

[0, 0, 1, 1, 2, 2]

nums = [2, 0, 1]

[0, 1, 2]

Решение

Задача Sort Colors также известна как "Задача о голландском национальном флаге". Это классика алгоритмических интервью.

Главная загвоздка: нельзя использовать встроенную сортировку и крайне желательно решить задачу за один проход O(n) с использованием фиксированного объема памяти O(1).

Хотя можно предложить метод сортировки подсчетом (Count Sort), на собеседованиях ждут именно вариант с одним проходом (алгоритм Дейкстры).

На собеседованиях решение через подсчет обычно считают уровнем для Junior или Junior+. Для Middle-разработчика это решение считается "базовым".

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

Однопроходный алгоритм позволяет проверить внимательность к деталям и умение работать с массивами "на месте". В реальных системах, например в низкоуровневой разработке, важно уметь менять данные в памяти. Кроме того, это база для более сложных задач: логика этой задачи используется в быстрой сортировке QuickSort при выборе опорного элемента и разделении массива на три части.

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

Если возникли сложности с решением, давайте подробнее его разберем.

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

Идея решения в том, что мы можем использовать два указателя: l и r, которые будут указывать нам на границы уже заполненных групп в начале и в конце массива для 0 и для 2.

Как только мы видим 0, ставим его в группу для 0 в начале массива. Если видим 2, то ставим его в группу для 2 в конце массива. Что делать с 1? Просто пропускаем.

Мы не просто копируем текущий элемент в l или r, а делаем операцию swap. Если видим 1, то идем дальше. Если видим 0, увеличиваем l и меняем текущий элемент с элементом на позиции l. Если видим 2, уменьшаем r и меняем текущий элемент с элементом на позиции r. Тут важно заметить, что в случае 2 мы делаем обмен и при этом не сдвигаем текущий элемент, так как мы должны проверить его на следующей итерации.

После прохода по массиву мы правильно заполним группы для 0 и 2, а также автоматически получим группу для 1 в середине массива.

Как и в задаче Сдвиг нулей, мы используем указатели как маркеры в памяти: "вот тут у меня все отсортировано, все, что дальше - еще нет".

Если подходить строго, можно сказать, что у нас три указателя - l, r и текущий индекс i, но основная идея решения строится именно на двух указателях l и r, которые движутся навстречу друг другу при заполнении групп.

Тут можно подумать над инвариантом цикла. Это поможет визуализировать решение, понять, почему оно верно и как в конце цикла мы получаем правильный ответ.

В этом алгоритме инвариант - это четкое распределение цветов по группам относительно указателей. Если в любой момент остановить цикл и проверить массив, то группы 0 и 2 должны находиться строго в своих границах, заданных указателями. Фактически у нас будут еще две группы: одна от l + 1 до i - 1 - здесь будут все 1, и одна от i до r - 1, где могут быть любые значения.

Итак, вот наш инвариант: в любой момент для текущего указателя i, левой границы l и правой границы r верно:

1. Все элементы в диапазоне [0, l] равны 0.
2. Все элементы в диапазоне [l + 1, i - 1] равны 1.
3. Все элементы в диапазоне [r, n - 1] равны 2.
4. Элементы в диапазоне [i, r - 1] могут быть любыми.

Начальные значения:

• l = -1 - нулей пока нет, группа 0 пуста.
• r = n - двоек пока нет, группа 2 пуста.
• i = 0 - начинаем проверку с первого элемента.

Вы можете самостоятельно попробовать доказать инвариант по методу математической индукции: проверьте начало и конец первой итерации, затем предположите, что инвариант верен на k - 1-й итерации, и проверьте начало и конец k-й итерации.

Приведем основные шаги алгоритма.

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

1. Инициализация: l = -1, r = n, i = 0.

2. Цикл: пока i < r.

   • Если текущий элемент nums[i] равен 0, увеличиваем l, меняем местами nums[l] и nums[i], затем увеличиваем i.
   • Если текущий элемент nums[i] равен 2, уменьшаем r и меняем местами nums[r] и nums[i].
   • Иначе текущий элемент равен 1, пропускаем его и увеличиваем i.

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

   • В итоге все 0 соберутся слева, 2 справа, а 1 в середине массива.

Вот как это работает на примере:

nums = [2, 0, 2, 1, 1, 0]

Начало: l = -1, r = 6, i = 0, nums = [2, 0, 2, 1, 1, 0]

1. nums[i], индекс 0, равен 2.
   • r--: стало 5, затем swap(nums[0], nums[5]).
   • Массив: [0, 0, 2, 1, 1, 2]
   • Итог: l = -1, r = 5, i = 0. Заметим: i остался на 0, чтобы проверить новый элемент.

2. nums[i], индекс 0, равен 0.
   • l++: стало 0, затем swap(nums[0], nums[0]) и i++: стало 1.
   • Массив: [0, 0, 2, 1, 1, 2]
   • Итог: l = 0, r = 5, i = 1.

3. nums[i], индекс 1, равен 0.
   • l++: стало 1, затем swap(nums[1], nums[1]) и i++: стало 2.
   • Массив: [0, 0, 2, 1, 1, 2]
   • Итог: l = 1, r = 5, i = 2.

4. nums[i], индекс 2, равен 2.
   • r--: стало 4, затем swap(nums[2], nums[4]).
   • Массив: [0, 0, 1, 1, 2, 2]
   • Итог: l = 1, r = 4, i = 2. Опять проверяем элемент по индексу 2.

5. nums[i], индекс 2, равен 1.
   • просто i++: стало 3.
   • Массив: [0, 0, 1, 1, 2, 2]
   • Итог: l = 1, r = 4, i = 3.

6. nums[i], индекс 3, равен 1.
   • просто i++: стало 4.
   • Массив: [0, 0, 1, 1, 2, 2]
   • Итог: l = 1, r = 4, i = 4.

7. Завершение
   • Условие цикла i < r, то есть 4 < 4, становится ложным.
   • Цикл окончен.

Итог: [0, 0, 1, 1, 2, 2]. Все нули слева от l = 1, все двойки справа от r = 4, единицы посередине.

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

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

Мы используем один цикл, который проходит по массиву от начала до конца ровно один раз. Таким образом, получаем временную сложность O(n).

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

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

Код решения

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

func sortColors(nums []int) {
    n := len(nums)
    l := -1
    r := n
    i := 0
    for i < r {
        if nums[i] == 0 {
            l++
            nums[i], nums[l] = nums[l], nums[i]
            i++
        } else if nums[i] == 2 {
            r--
            nums[i], nums[r] = nums[r], nums[i]
        } else {
            i++
        }
    }
}

Итоги

Вот основные итоги по задаче, которые полезно держать в голове для собеседования:

1. Суть алгоритма. Мы использовали разделение массива на группы для 0 и 2 с помощью двух указателей.

2. Эффективность. В отличие от метода подсчета, который требует двух проходов, это решение работает за один проход.

3. Подводные камни:

   • Инициализация l = -1 и r = n упрощает логику, делая группы пустыми до начала цикла.
   • Не увеличивать i при обмене с двойкой. Это критически важно, так как из конца массива может прийти любой цвет, который еще нужно проверить.
   • Условие остановки цикла - i < r.

4. Где пригодится. Эта техника разделения массива на части, partitioning, лежит в основе:

   • Алгоритма QuickSort, или быстрой сортировки.
   • Задачи Kth Largest Element через QuickSelect.
   • Задач на группировку элементов "на месте" по какому-то признаку.