---

Условие задачи

Дан массив, содержащий только 0, 1 и 2. Отсортируйте его на месте. Считайте элементы объектами: нельзя подсчитать количество значений каждого вида и заново заполнить массив.

Задача получила название по флагу Нидерландов, который состоит из трёх цветов.

Пример 1:

Вход: [1, 0, 2, 1, 0]
Выход: [0, 0, 1, 1, 2]

Пример 2:

Вход: [2, 2, 0, 1, 2, 0]
Выход: [0, 0, 1, 2, 2, 2]

Решение

Сортировка на месте, например пирамидальная, потребовала бы \(O(N \log N)\) времени. Здесь достаточно одного прохода с границами low и high:

• перед low находятся только нули;
• после high находятся только двойки;
• между low и текущим индексом находятся единицы.

При встрече нуля меняем его местами с элементом на low и сдвигаем оба индекса. Единицу просто пропускаем. Двойку меняем местами с элементом на high и уменьшаем только high, потому что новый элемент на текущей позиции ещё не проверен.

Код

package main

import "fmt"

func sortDutchFlag(arr []int) {
	low, high := 0, len(arr)-1
	for i := 0; i <= high; {
		switch arr[i] {
		case 0:
			arr[i], arr[low] = arr[low], arr[i]
			i++
			low++
		case 1:
			i++
		case 2:
			arr[i], arr[high] = arr[high], arr[i]
			high--
		}
	}
}

func main() {
	first := []int{1, 0, 2, 1, 0}
	sortDutchFlag(first)
	fmt.Println(first)

	second := []int{2, 2, 0, 1, 2, 0}
	sortDutchFlag(second)
	fmt.Println(second)
}

Вывод:

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

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

Каждый элемент проверяется не более одного раза, поэтому временная сложность равна \(O(N)\).

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

Алгоритм сортирует массив на месте и использует \(O(1)\) дополнительной памяти.

---