Избегайте вложенных циклов

19.09.2021

Бывают задачи, в которых вам надо пройтись по какой-то одномерной структуре (массиву, строке), но при этом в определенных случаях вам может захотеться писать вложенные циклы — как правило, двигающие одну и ту же переменную.

Простейший пример: пусть вам надо в строке заменить все группы подряд идущих пробелов на один, т.е. из строки ab cd e сделать ab cd e. Вы можете попробовать написать что-нибудь следующего рода:

i = 0
res = ""
while i < len(s):
    if s[i] != " ":
        res += s[i]
        i += 1
    else:
        res += " "
        while s[i] == " ":
            i += 1

Казалось бы логично: все не-пробельные символы копируем, а из пробельных копируем один, а остальные пропускаем. Но на самом деле такого рода код обычно получается сравнительно сложным и создает риск ошибок; в частности, в коде выше есть ошибка (не очень сложная, но попробуйте ее найти; ответ будет дальше по тексту).

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

for i in range(len(s)):
    if s[i] != " ":
        res += s[i]
    else:
        if i == 0 or s[i - 1] != " ":
            res += s[i]

Смысл простой: мы идем по строке, и про каждый символ думаем, надо его копировать в выходную строку или нет. Копировать надо, если это не пробел, а также если это первый пробел в группе последовательных пробелов.

В каком-то смысле вместо внутреннего цикла мы просто делаем один, первый, шаг этого цикла — ну и не страшно, на следующей итерации внешнего цикла мы сделаем следующий шаг, и т.д.

Код получается заметно проще, по крайней мере по следующим параметрам. Во-первых, внешний цикл получился for, а не while, и мы нигде не меняет вручную переменную i; это избавляет нас от риска того, что мы забудем написать i+=1 в какой-нибудь из веток. Во-вторых, нет риска выхода за пределы строки: как только i дойдет до конца строки, цикл автоматом закончится. А в коде выше, с вложенным циклом, ошибка как раз в выходе за пределы строки во внутреннем цикле — и это на самом деле очень типичная и характерная ошибка в таких задачах. В-третьих, за счет того, что код на каждой итерации работает только с одним символом, становится намного проще его понимать и искать ошибки. За счет того, что на каждой итерации вы делаете только один шаг, намного проще понимать, в каком порядке меняются ваши переменные цикла, и проще понимать, что в целом происходит.

Еще один яркий пример. Пусть нам дана строка, и надо проверить, явлется ли она палиндромом, если удалить все пробелы (например, классическое «а роза упала на лапу азора» является палиндромом в этом смысле). Самый простой, конечно, способ — скопировать все не-пробельные символы во вторую строку, и дальше ее простым образом проверить, но давайте попробуем обойтись без дополнительной строки. Казалось бы, логично, запустим два указателя с двух концов строки, будем ими пропускать пробелы, а не-пробельные символы сравнивать, что-то типа такого:

def check(s):
    i = 0
    j = len(s)
    while i < j:
        while i < len(s) and s[i] == " ":
            i += 1
        while j >= 0 and s[j] == " ":
            j -= 1
        if s[i] != s[j]:
            return False

Но тут тоже есть ошибки. Во-первых, забыт сдвиг i += 1 и j -= 1 в конце цикла, во-вторых, если строка состоит из одних только пробелов, то в if будет выход за пределы строки, ну и, возможно, есть какие-то еще проблемы, которые я сразу не заметил.

(Плюс ко всему, этот код делает одно лишнее сравнение, если в середине строки есть пробелы — тогда i и j успеют поменяться местами, но мы s[i] и s[j] будем все равно сравнивать. В этой задаче это не создаст проблем, кроме как в упомянутом выше случае строки из одних пробелов, но в других ситуациях это может создать свои проблемы.)

Но можно написать так:

def check(s):
    i = 0
    j = len(s)
    while i < j:
        if s[i] == " ":
            i += 1
            continue
        if s[j] == " ":
            j -= 1
            continue
        if s[i] != s[j]:
            return False
        i += 1
        j -= 1

и все становится проще.

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

Еще пример — любые задачи, решающиеся классическим подходом с созданием конечного автомата. Простейший пример — задача Смайлики. В ней тоже можно написать код с двумя вложенными циклами: во внешнем цикле ищем начало смайлика, а во внутреннем ищем конец смайлика, проверяя символы по пути. Но можно ввести достаточно формализованное понятие «состояния» — в данном случае можно выделить четыре состояния: вне смайлика, внутри «глаз» смайлика, внутри «носа» и внутри «рта», и несложно написать правила, которые по символу и текущему состоянию определят следующий. Это чистой воды классический конечный автомат; вообще, описанные выше примеры тоже очень похожи на подход с конечными автоматами (особенно мои рассуждения, что намного проще понимать, что происходит за один ход), но все-таки в общем случае это далеко не обязательно строго формализованный конечный автомат.

Поэтому, когда вы хотите написать два вложенных цикла, двигающих одну и ту же переменную — подумайте, может быть, можно обойтись одним циклом, обрабатывая только один объект за раз?


Мой курс по алгоритмическому программированию (и подготовке к олимпиадам) для школьников, студентов и всех желающих — algoprog.ru