Rotate Array: циклический сдвиг массива на k за O(n) и O(1) памяти

Коротко

Задача выглядит простой, пока не попросят решить её без лишней памяти. Тут и начинается интересное: три совершенно разных подхода, немного теории чисел и пара ловушек, на которых валятся даже аккуратные кандидаты. Разберём их по порядку.

Условие

Дан массив nums и неотрицательное число k. Нужно циклически сдвинуть массив вправо на k позиций. Элементы, которые «вываливаются» справа, появляются слева.

Вход:  nums = [1, 2, 3, 4, 5, 6, 7], k = 3
Выход: [5, 6, 7, 1, 2, 3, 4]

Вход:  nums = [-1, -100, 3, 99], k = 2
Выход: [3, 99, -1, -100]

На LeetCode это задача 189, функция меняет массив на месте и ничего не возвращает. В примерах ниже мы для удобства возвращаем nums.

Главная формула, на которой стоит всё остальное: элемент с индексом i должен оказаться на позиции (i + k) % n. Обратное направление тоже полезно: на позицию j приезжает элемент, который стоял на (j - k) % n.

Сначала про k > n

Ограничения в задаче обычно такие: n до $10^5$, а k до $10^5$ или даже больше. То есть k вполне может превышать длину массива. Это первое, что проверяет интервьюер.

Сдвиг на n позиций ничего не меняет: каждый элемент делает полный круг и возвращается на место. Значит, сдвиг на k равен сдвигу на k % n. Для [1, 2, 3] и k = 7 реальный сдвиг равен 7 % 3 = 1, ответ [3, 1, 2].

Если забыть про остаток, решения ломаются по-разному: одно падает с выходом за границы, другое молча возвращает неверный ответ, третье тратит $O(n \cdot k)$ времени впустую. Примеры будут ниже. Правило простое: k %= n первой строкой, а перед этим проверка на пустой массив, иначе получите деление на ноль.

В Python остаток от отрицательного числа всегда неотрицательный: -2 % 5 даёт 3. В JavaScript, Java и C++ знак остатка совпадает со знаком делимого, и -2 % 5 равно -2. Если по условию k может быть отрицательным (сдвиг влево), нормализуйте так: ((k % n) + n) % n. Подробнее про остатки — в статье о модульной арифметике.

Наивное решение: сдвигать по одному

Самое очевидное: k раз сдвинуть массив на одну позицию вправо. Запоминаем последний элемент, остальные сдвигаем на шаг, запомненный ставим в начало.

def rotate_by_one(nums, k):
    n = len(nums)
    if n == 0:
        return nums
    for _ in range(k % n):
        last = nums[-1]
        for i in range(n - 1, 0, -1):
            nums[i] = nums[i - 1]
        nums[0] = last
    return nums

Памяти $O(1)$, но времени $O(n \cdot k)$. При $n = k = 10^5$ это $10^{10}$ операций — тайм-аут гарантирован. Без k % n всё ещё хуже: при k = 10^9 цикл крутится миллиард раз ради результата, который совпадает со сдвигом на k % n.

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

Решение 1: дополнительный массив

Раз мы знаем, куда должен попасть каждый элемент, просто положим его туда в новом массиве.

def rotate_copy(nums, k):
    n = len(nums)
    if n == 0:
        return nums
    result = [0] * n
    for i in range(n):
        result[(i + k) % n] = nums[i]
    nums[:] = result
    return nums

Здесь % n стоит внутри цикла, поэтому большое k не страшно. Время $O(n)$, память $O(n)$.

Строка nums[:] = result принципиальна: она копирует значения в существующий список. Если написать nums = result, поменяется только локальное имя, а массив вызывающего кода останется прежним. Это та же ловушка, что и в Move Zeroes.

В Python то же самое пишется через срезы:

def rotate_slices(nums, k):
    if not nums:
        return nums
    k %= len(nums)
    nums[:] = nums[-k:] + nums[:-k]
    return nums

Последние k элементов уходят в начало, остальные — за ними. Обратите внимание на k %= len(nums). Без неё при k = 7 и [1, 2, 3] срез nums[-7:] вернёт весь массив, nums[:-7] — пустой, и функция молча вернёт исходный [1, 2, 3] вместо [3, 1, 2]. Ошибки нет, ответ неверный — худший вид бага.

Случай k = 0 здесь, кстати, работает: nums[-0:] это nums[0:], весь массив, а nums[:0] пустой. Но полагаться на такие совпадения не стоит, лучше проговорить его вслух.

Для продакшен-кода на Python есть ещё collections.deque с методом rotate(k): положительное k сдвигает вправо, отрицательное — влево. На собеседовании его можно упомянуть, но решать задачу через него не дадут. Про сам дек — в отдельной статье.

Обычно после этого решения звучит вопрос: «А можно за $O(1)$ памяти?» Можно, и двумя способами.

Решение 2: жонглирование циклами

Идея

Возьмём элемент на позиции 0. Он должен переехать на позицию k. Тот, кто стоял на k, должен переехать на 2k, и так далее, по модулю n. Получается цепочка: берём элемент, ставим его на место, вытесненного несём дальше. Рано или поздно мы вернёмся на позицию 0, и цикл замкнётся.

Для n = 7, k = 3 цепочка с нуля обходит все позиции: 0 → 3 → 6 → 2 → 5 → 1 → 4 → 0. Один цикл, и массив готов.

А для n = 6, k = 2 цепочка из нуля такая: 0 → 2 → 4 → 0. Позиции 1, 3, 5 остались нетронутыми. Нужен второй цикл с позиции 1: 1 → 3 → 5 → 1.

Сколько будет циклов

Начиная с позиции s, мы посещаем позиции $s + jk \bmod n$ для $j = 0, 1, 2, \ldots$ Возврат в s случится при наименьшем $j$, для которого $jk$ делится на $n$. Это $j = n / \gcd(n, k)$.

Значит, каждый цикл имеет длину $n / g$, где $g = \gcd(n, k)$, а всего циклов $g$. Ещё одно наблюдение: $k$ делится на $g$, поэтому все позиции одного цикла дают одинаковый остаток при делении на $g$. Отсюда стартовые точки — это просто $0, 1, \ldots, g - 1$: каждая лежит в своём цикле.

Проверим на примерах:

Отсюда и название: алгоритм жонглирования (juggling algorithm), а НОД считается алгоритмом Евклида за $O(\log n)$.

Код

from math import gcd


def rotate_cycles(nums, k):
    n = len(nums)
    if n == 0:
        return nums
    k %= n
    if k == 0:
        return nums
    for start in range(gcd(n, k)):
        current = start
        value = nums[start]
        while True:
            nxt = (current + k) % n
            nums[nxt], value = value, nums[nxt]
            current = nxt
            if current == start:
                break
    return nums

В value лежит элемент, который мы «несём в руках». На каждом шаге кладём его на позицию nxt и забираем то, что там было. Когда вернулись в start, кладём последний элемент на исходную позицию цикла, и он замкнут.

Каждый элемент записывается ровно один раз — всего $n$ записей. Время $O(n)$, память $O(1)$.

Вариант без НОД

Если не хочется доказывать формулу на собеседовании, можно просто считать перемещённые элементы и запускать новый цикл со следующей позиции, пока не наберётся n:

def rotate_cycles_count(nums, k):
    n = len(nums)
    if n == 0:
        return nums
    k %= n
    moved = 0
    start = 0
    while moved < n:
        current = start
        value = nums[start]
        while True:
            nxt = (current + k) % n
            nums[nxt], value = value, nums[nxt]
            current = nxt
            moved += 1
            if current == start:
                break
        start += 1
    return nums

Почему start += 1 не попадает в уже обработанный цикл? Ровно по причине выше: стартовые точки $0, \ldots, g - 1$ лежат в разных циклах, а к моменту, когда start дойдёт до g, счётчик уже наберёт n. Здесь k = 0 не требует отдельной проверки: каждый «цикл» имеет длину 1, элемент записывается сам в себя.

Такой вариант часто встречается в разборах задачи. Он рабочий, но кандидаты часто путаются в нём под стрессом: забывают обновить value или проверяют условие выхода не в том месте. Поэтому на практике чаще пишут следующее решение.

Решение 3: три разворота

Идея

Посмотрим на массив как на два куска: A — первые n - k элементов, B — последние k. Сдвиг вправо на k превращает AB в BA.

Обозначим через $X^R$ развёрнутый кусок $X$. Есть простое тождество:

$(AB)^R = B^R A^R$

Развернув весь массив, мы уже поставили B перед A, только оба куска лежат задом наперёд. Осталось развернуть каждый на своём месте: $(B^R)^R = B$, $(A^R)^R = A$. Получаем BA.

Разбор на примере

nums = [1, 2, 3, 4, 5, 6, 7], k = 3. Здесь A = [1, 2, 3, 4], B = [5, 6, 7].

Код

def reverse(nums, left, right):
    while left < right:
        nums[left], nums[right] = nums[right], nums[left]
        left += 1
        right -= 1


def rotate_reverse(nums, k):
    n = len(nums)
    if n == 0:
        return nums
    k %= n
    reverse(nums, 0, n - 1)
    reverse(nums, 0, k - 1)
    reverse(nums, k, n - 1)
    return nums


print(rotate_reverse([1, 2, 3, 4, 5, 6, 7], 3))  # [5, 6, 7, 1, 2, 3, 4]
print(rotate_reverse([-1, -100, 3, 99], 2))  # [3, 99, -1, -100]
print(rotate_reverse([1, 2, 3], 7))  # [3, 1, 2]

Функция reverse — классические встречные два указателя. Крайние случаи обрабатываются сами: при k = 0 вызов reverse(nums, 0, -1) ничего не делает, потому что left < right ложно сразу, и массив разворачивается дважды целиком — то есть остаётся прежним.

А вот без k %= n здесь будет реальная беда. При k = 7 и n = 3 второй вызов превратится в reverse(nums, 0, 6), и Python упадёт с IndexError. В C++ это и вовсе неопределённое поведение.

Сложность

Три разворота делают в сумме примерно $n$ обменов: $\lfloor n/2 \rfloor + \lfloor k/2 \rfloor + \lfloor (n-k)/2 \rfloor \le n$. Время $O(n)$, память $O(1)$. Каждый элемент перемещается дважды, а не один раз, как в жонглировании, но код короче, а ошибиться в нём почти негде. Освежить, как считается асимптотика, можно в разборе O-нотации.

То же на JavaScript

function reverse(nums, left, right) {
  while (left < right) {
    ;[nums[left], nums[right]] = [nums[right], nums[left]]
    left++
    right--
  }
}

function rotate(nums, k) {
  const n = nums.length
  if (n === 0) return nums
  k = ((k % n) + n) % n
  reverse(nums, 0, n - 1)
  reverse(nums, 0, k - 1)
  reverse(nums, k, n - 1)
  return nums
}

console.log(rotate([1, 2, 3, 4, 5, 6, 7], 3)) // [5, 6, 7, 1, 2, 3, 4]
console.log(rotate([1, 2, 3, 4, 5], -2)) // [3, 4, 5, 1, 2]

Нормализация ((k % n) + n) % n здесь позволяет передать отрицательное k и получить сдвиг влево. Без + n при k = -2 вызов reverse(nums, 0, -3) ничего бы не сделал, а reverse(nums, -2, 4) полез бы по отрицательным индексам.

Сдвиг влево

Сдвиг влево на k — это сдвиг вправо на n - k % n. Можно переиспользовать готовую функцию:

def rotate_left(nums, k):
    n = len(nums)
    if n == 0:
        return nums
    return rotate_reverse(nums, n - k % n)


print(rotate_left([1, 2, 3, 4, 5], 2))  # [3, 4, 5, 1, 2]

Если k % n == 0, получится сдвиг вправо на n, а rotate_reverse сама сведёт его к нулю.

Есть и прямой вариант тремя разворотами: сначала развернуть первые k, потом остальные, потом весь массив. Это то же тождество, только в другом порядке: $(A^R B^R)^R = BA$, где теперь A — первые k элементов.

def rotate_left_reverse(nums, k):
    n = len(nums)
    if n == 0:
        return nums
    k %= n
    reverse(nums, 0, k - 1)
    reverse(nums, k, n - 1)
    reverse(nums, 0, n - 1)
    return nums

Именно в таком виде этот приём описан в книге Джона Бентли «Жемчужины программирования», где сравниваются все три способа поворота на месте.

Сравнение подходов

Если ограничений по памяти нет, дополнительный массив — нормальный выбор, его проще всего прочитать. Если нужно на месте, пишите три разворота. Жонглирование хорошо знать для разговора: оно делает минимум записей и показывает, что вы понимаете структуру перестановки.

Любопытная деталь: в реальном железе жонглирование не всегда быстрее разворотов, несмотря на вдвое меньшее число записей. Оно прыгает по памяти с шагом k, а развороты идут последовательно и лучше дружат с кешем процессора. Так что «меньше операций» и «быстрее» — не одно и то же.

Типичные ошибки

Забыть k %= n. Самая частая. Последствия зависят от решения: IndexError в разворотах, неверный ответ в срезах, лишняя работа в наивном варианте.

Не проверить пустой массив. k %= 0 бросает ZeroDivisionError. На LeetCode n >= 1, но интервьюер может спросить.

nums = result вместо nums[:] = result. Решение «работает» в вашей голове и не меняет ничего снаружи.

Перепутать направление. Вправо на k: последние k элементов уходят в начало. Влево: первые k уходят в конец. Прогоните маленький пример на бумаге до того, как писать код.

Перепутать границы второго и третьего разворота. Правильно [0, k - 1] и [k, n - 1]. Вариант [0, k] и [k + 1, n - 1] даёт сдвиг на k + 1.

Один цикл в жонглировании. Решение, которое проходит одну цепочку с позиции 0, работает только при $\gcd(n, k) = 1$. На n = 7, k = 3 оно пройдёт, на n = 6, k = 2 — нет. Поэтому тестируйте на чётных n и чётных k.

Отрицательный остаток в JS/Java/C++. Если k может быть отрицательным, k % n тоже окажется отрицательным. Нормализуйте через ((k % n) + n) % n.

Где ещё встречается

Циклический сдвиг — базовая операция, и он всплывает во многих задачах:

• Поиск в повёрнутом отсортированном массиве. Отсортированный массив сдвинули на неизвестное k, нужно найти элемент или минимум за $O(\log n)$. Разборы — в статьях про поиск в повёрнутом массиве и поиск минимума. Найденный индекс минимума и есть величина сдвига.
• Повернуть строку. Задача «является ли строка goal сдвигом строки s» решается проверкой len(s) == len(goal) and goal in s + s: все сдвиги s являются подстроками s + s.
• Кольцевой буфер. Вместо физического сдвига элементов хранят индекс начала и обращаются по (head + i) % n. Тот же модуль, только без перемещения данных. Подробнее — в статье про кольцевой буфер.
• Rotate List. Тот же сдвиг для связного списка: там разворачивать дорого, зато можно замкнуть список в кольцо и разрезать в нужном месте.
• Поворот матрицы на 90°. Похожий приём с разворотами: транспонировать и развернуть каждую строку.

На собеседовании

Rotate Array — частая задача уровня medium на массивы. Её любят как раз за то, что «правильных» решений несколько, и по ходу разговора видно, как кандидат рассуждает. Хороший порядок ответа:

1. Уточните условие: направление сдвига, может ли k быть больше n или отрицательным, нужно ли на месте.
2. Сразу скажите про k %= n и почему это корректно.
3. Назовите решение с дополнительным массивом и его $O(n)$ по памяти.
4. Напишите три разворота и объясните тождество $(AB)^R = B^R A^R$ на примере.
5. Прогоните крайние случаи: k = 0, k = n, k > n, массив из одного элемента.

Чего ждать в продолжении:

• «Можно ли сделать меньше записей?» — расскажите про жонглирование и $\gcd(n, k)$ циклов.
• «Почему циклов именно столько?» — длина цикла $n / \gcd$, потому что это наименьшее $j$, при котором $jk$ делится на $n$.
• «А сдвиг влево?» — сдвиг вправо на n - k % n или развороты в другом порядке.
• «А если это связный список?» — замкнуть в кольцо и разрезать.

Больше типовых задач этого уровня собрано в обзоре задач на массивы на собеседовании. Чтобы закрепить приём, напишите три разворота с нуля без подсказок, а потом попробуйте жонглирование и проверьте его на n = 6, k = 4. Если оба проходят с первого раза, задача усвоена.