Отсортируй массив светофора в порядке r-y-g

Такое дают на разминку в начале алгоритмической секции: массив, где всего три вида элементов, надо разложить в фиксированном порядке. Хочется сразу написать sorted(), но собеседующий ждёт, что ты заметишь — раз значений всего три, обычная сортировка избыточна.

Задача

Дан массив из символов 'r', 'y', 'g' (красный, жёлтый, зелёный — цвета светофора) в произвольном порядке. Нужно вернуть массив, отсортированный в порядке светофора: сначала все 'r', потом все 'y', потом все 'g'.

Например, ['g', 'g', 'y', 'y', 'r', 'r'] → ['r', 'r', 'y', 'y', 'g', 'g'].

Решение

def sort_traffic(arr, order=("r", "y", "g")):
    result = []
    for color in order:
        result += color * arr.count(color)
    return result


cases = [
    ["g", "g", "y", "y", "r", "r"],
    ["g", "r", "g", "y", "r"],
    ["y", "y", "y"],
]
for arr in cases:
    print(f"{arr} -> {sort_traffic(arr)}")

Разбор

Главная мысль: когда различных значений всего несколько (тут три), не надо ничего сравнивать попарно — достаточно посчитать, сколько раз встретился каждый, и выписать их в нужном порядке. Это сортировка подсчётом (counting sort).

В решении мы задаём порядок как кортеж ('r', 'y', 'g') и идём по нему. Для каждого цвета считаем arr.count(color) — сколько таких в массиве — и дописываем в результат столько же копий. Приём color * n для строки повторяет символ n раз ('r' * 2 → 'rr'), а result += 'rr' расширяет список посимвольно. В итоге результат собирается блоками: сначала весь красный блок, затем жёлтый, затем зелёный.

Почему это лучше sorted(). Обычная сортировка сравнениями работает за O(n log n). Здесь же мы делаем фиксированное число проходов по массиву (по одному count на каждый из трёх цветов) — это O(n), линейно. Множитель — число различных значений, а оно константно (три), поэтому в асимптотике не участвует. На больших массивах разница ощутима.

Есть нюанс, который стоит проговорить: в этом решении мы делаем три полных прохода (по одному count на цвет) и строим новый массив. Классический алгоритм голландского флага (Dutch national flag problem, придуман Дейкстрой ровно про три «цвета») решает ту же задачу за один проход и на месте, тремя указателями, вообще без выделения новой памяти. Для собеседования на джуна вариант с подсчётом — отличный и понятный ответ; упоминание, что «в пределе это задача о голландском флаге и её можно сделать in-place за один проход», покажет глубину.

Порядок задаётся параметром, так что решение легко обобщается на любой набор ключей и любой порядок сортировки. Родственная задача на сборку массива из кусочков — чередование двух массивов.

Ожидаемый результат

['g', 'g', 'y', 'y', 'r', 'r'] -> ['r', 'r', 'y', 'y', 'g', 'g']
['g', 'r', 'g', 'y', 'r'] -> ['r', 'r', 'y', 'g', 'g']
['y', 'y', 'y'] -> ['y', 'y', 'y']