Слей два массива с чередованием элементов

Такое дают на скрининге как разминочную задачу: два массива надо «переплести» в один, беря по очереди из первого и второго. Пишется в несколько строк, но проверяет, уверенно ли ты работаешь с индексами и думаешь ли про длины массивов.

Задача

Даны два массива одинаковой длины, например a = [1, 8, 9] и b = [2, 3, 1]. Нужно собрать один массив, чередуя элементы: сначала первый элемент a, затем первый элемент b, затем второй a, второй b, и так далее.

Для примера выше результат — [1, 2, 8, 3, 9, 1].

Решение

def interleave(a, b):
    result = []
    for i in range(len(a)):
        result.append(a[i])
        result.append(b[i])
    return result


cases = [
    ([1, 8, 9], [2, 3, 1]),
    ([1, 2, 3], [4, 5, 6]),
]
for a, b in cases:
    print(f"{a} + {b} -> {interleave(a, b)}")

Разбор

Идея прямолинейная: идём по индексам от 0 до len(a) - 1 и на каждом шаге кладём в результат сначала элемент из a, потом из b с тем же индексом. За один проход собирается переплетённый массив.

Почему через индекс, а не через два цикла подряд. Если сначала выгрузить весь a, а потом весь b, получится конкатенация [1, 8, 9, 2, 3, 1] — не то. Чередование требует синхронного шага по обоим массивам, и индекс i — самый явный способ этот шаг выразить: он один и тот же для a[i] и b[i].

Идиоматичная альтернатива в Python — zip, который как раз ходит по двум последовательностям параллельно:

def interleave(a, b):
    result = []
    for x, y in zip(a, b):
        result.append(x)
        result.append(y)
    return result

Читается чище и не требует ручной работы с индексами. Хороший ход на собеседовании — написать индексный вариант, а потом предложить zip как более «питоничный».

Главная ловушка — разная длина. Индексный цикл идёт по len(a), поэтому если b короче — вылетит IndexError при обращении к несуществующему b[i]; если b длиннее — его хвост молча потеряется. zip ведёт себя иначе: он останавливается на более коротком массиве и лишнее просто отбрасывает без ошибки. Стоит уточнить у собеседующего, гарантированы ли равные длины, и если нет — обсудить, что делать с хвостом (дописать остаток через itertools.zip_longest или обрезать). Именно вопрос про краевой случай отличает вдумчивое решение от «на глаз работает».

Сложность — O(n): один проход, каждый элемент трогаем ровно раз. Родственная задача на сборку результата по правилу — сортировка светофора.

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

[1, 8, 9] + [2, 3, 1] -> [1, 2, 8, 3, 9, 1]
[1, 2, 3] + [4, 5, 6] -> [1, 4, 2, 5, 3, 6]