Задача с собеседования в Google: Sort Colors

#algorithms #google #interview #leetcode

Задача

Дан массив целых чисел. Значения могут быть только 0, 1 и 2. Нужно отсортировать на месте (без доп. памяти). Использовать библиотечные функции сортировки нельзя.

Например,

Input: nums = [2,0,2,1,1,0]
Output: [0,0,1,1,2,2]

Input: nums = [2,0,1]
Output: [0,1,2]

Ссылка на leetcode: https://leetcode.com/problems/sort-colors

Решение

Сортировки вроде QuickSort, MergeSort не подходят, т.к. используют дополнительную память. QuickSort - O(log(n)) для рекурсии, MergeSort - O(n) для временного массива.

Простые сортировки, вроде Bubble Sort, Insertion Sort, Selection Sort дополнительной памяти не требуют и работают за O(1), но временная сложность - O(n^2). Хотелось бы побыстрее.

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

В коде это выглядит так:

int zeros = 0, ones = 0, twos = 0;

//считаем, сколько у нас нулей, единичек и двоек
for (int num : nums) {
    if (num == 0) {
        zeros++;
    } else if (num == 1) {
        ones++;
    } else {
        twos++;
    }
}

//записываем нужное число нулей, единичек и двоек в порядке возрастания
int i = 0;
while (zeros > 0) {
    nums[i++] = 0;
    zeros--;
}
while (ones > 0) {
    nums[i++] = 1;
    ones--;
}
while (twos > 0) {
    nums[i++] = 2;
    twos--;
}

Time Complexity = O(2*n) = O(n)
Space Complexity = O(1)

Работает без доп. памяти за линейное время.

Можно ли улучшить это решение? Да, можно. Можно сделать также за линейное время без доп. памяти, но за один проход, а не за два.

Dutch national flag problem

Вообще, это известная задача - Dutch national flag problem. Ее предложил Edsger W. Dijkstra. Он сам родом из Нидерландов.

Эту задачу можно решить за один проход при помощи трех указателей (Three pointers).

Для этого нужно объявить 3 указателя, которые делят массив на 4 области:
1) [0..left-1] — все нули
2) [left..current-1] — все единички
3) [current..right] - еще необработанная зона
4) [right+1..N-1] - все двойки

left - определяет границу, слева от которой, у нас только нули.
right - определяет границу, справа от которой, у нас только двойки.
current - текущий элемент, опеределяет две границы: 1) слева до left включительно - только единички 2) справа до right включительно - еще не известно (необработанная зона).

Например:

Зелёным цветом я выделил первую область — там, где только нолики; синим — вторую, где только единички; красным — необработанную область; и, наконец, оранжевым — ту, где только двойки.

Нам нужно выбрать такую стратегию изменения указателей, чтобы эти правила (инварианты) соблюдались всегда во время работы алгоритма.

Изменять указатели будем до тех пор, пока current не станет больше right (current > right). В таком случае у нас исчезнет область 3 (необработанная), т.к. она [current..right], а в конце у нас current > right. Т.е. эта область станет пустой. Тогда у нас останутся три области, которые будут идти одна за другой в нужном нам порядке.

Инициализируем указатели следующими значениями:
left = 0
current = 0
right = N-1

Проверим, соблюдаются ли в начале наши инварианты:
1) [0..left-1] — все нули -> [0..-1] - пустая область, соблюдается
2) [left..current-1] — все единички -> [0..-1] -> пустая область, соблюдается
3) [current..right] - еще необработанная зона -> [0..N-1] - весь массив необработан, соблюдается
4) [right+1..N-1] - все двойки -> [N..N-1] - пустая область, соблюдается

При проходе массива мы можем встретить 3 значения (0, 1, 2). Давайте посмотрим, что делать при каждом из трех вариантов, чтобы сохранить наши инварианты (области):
1) nums[current] == 0.
Меняем местами nums[left] и nums[current], инкрементируем left и current. Тогда nums[left] станет равным 0 и можно расширить область с ноликами (left++). На nums[current] станет то, что было на nums[left]. А что там могло быть?

Если мы соблюдали инварианты - там могла быть единичка. Если единичка, то она встанет на место nums[current] = 1 и область с единичками можно расширить. Поэтому current++. Необработанная область уменьшится на 1, а область с двойками не изменится. Так, что в этом случае все инварианты соблюдаются.
Там мог быть нолик. Если область с единичками у нас пустая. left == current. Тогда swap, по факту, ничего не переставит. Мы увеличим область с ноликами, область с единичками останется пустой left == current.
Двойки там быть не могло, для этого left должно быть больше current.

2) nums[current] == 1. Делаем только current++. Области с ноликами и двойками не меняются. Область с единичками растет на 1 и область с необработанными числами сокращается на 1.

3) nums[current] == 2. Меняем местами nums[right] и nums[current], уменьшаем на единицу right: right--.
Это увеличит область с двойками на единицу. Остальные области не изменятся. Почему по аналогии нельзя увеличить current? В данном случае, nums[right] был в необработанной области и мы не знаем какое там было значение. Поэтому nums[current] после обмена останется неизвестным с необработанной зоне.

Почему мы выйдем из цикла? на каждой итерации мы увелициваем или current или уменьшаем right. Рано или поздно, current станет больше, чем right.

В коде это выглядит так:

public void sortColors(int[] nums) {
    int left = 0;
    int right = nums.length - 1;
    int current = 0;
    while (current <= right) {
        if (nums[current] == 0) {
            swap(nums, left, current);
            left++;
            current++;
        } else if (nums[current] == 1) {
            current++;
        } else {
            swap(nums, current, right);
            right--;
        }
    }
}

private void swap(int nums[], int i, int j) {
    int tmp = nums[i];
    nums[i] = nums[j];
    nums[j] = tmp;
}