Итератор раскрывающий многоуровневые списки

Рейтинг: 2Ответов: 2Опубликовано: 29.03.2023

Необходимо написать итератор раскрывающий многоуровневые списки вида:

list_of_lists_2 = [
    [['a'], ['b', 'c']],
    ['d', 'e', [['f'], 'h'], False],
    [1, 2, None, [[[[['!']]]]], []]
]

Написал генератор выполняющий данную задачу:


def flat_generator(list_of_list):
    for item in list_of_list:
        if type(item) is not list:
            yield item
        else:
            yield from flat_generator(item)

По тому же принципу (с рекурсией) пишу итератор:

class FlatIterator:

    def __init__(self, list_of_list):
        self.list_of_list = list_of_list

    def __iter__(self):
        self.counter = -1
        return self

    def __next__(self):

        self.counter += 1
        if self.counter == len(self.list_of_list):
            raise StopIteration
        item = self.list_of_list[self.counter]
        if type(item) is not list:
            return item
        else:
            for new_item in FlatIterator(item):
               return new_item ?????

Но второй return возвращает только первый элемент списка, при использовании yield получаю бесконечный цикл( Как мне исправить мой код?

python итераторы
Источник: Stack Overflow на русском

Ответы

▲ 3Принят

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

class FlatIterator:
    def __init__(self, list_of_list):
        self.stack = []
        self.current_list = list_of_list
        self.counter = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.counter >= len(self.current_list):
            if self.stack:
                # Дошли до конца вложенного списка,
                # выходим из него во внешний, "вспоминаем" состояние
                self.current_list, self.counter = self.stack.pop()
                return next(self)
            else:
                raise StopIteration
        
        item = self.current_list[self.counter]
        self.counter += 1

        if type(item) is not list:
            return item
        else:
            # Запоминаем текущее состояние
            self.stack.append((self.current_list, self.counter))
            # Входим во вложенный список, он становится текущим списком
            self.current_list = item
            self.counter = 0
            return next(self)


list_of_lists_2 = [
    [['a'], ['b', 'c']],
    ['d', 'e', [['f'], 'h'], False],
    [1, 2, None, [[[[['!']]]]], []],
]

print(list(FlatIterator(list_of_lists_2)))

Вывод:

['a', 'b', 'c', 'd', 'e', 'f', 'h', False, 1, 2, None, '!']

В общем, то еще "удовольствие", через рекурсивную функцию-генератор получается намного проще.

Можно внутри класса использовать итераторы вместо списков, тогда не нужно будет хранить индекс в списке:

class FlatIterator:
    def __init__(self, list_of_list):
        self.stack = []
        self.current_iterator = iter(list_of_list)

    def __iter__(self):
        return self

    def __next__(self):
        try:
            item = next(self.current_iterator)
        except StopIteration:
            if self.stack:
                # Дошли до конца вложенного списка,
                # выходим из него во внешний, "вспоминаем" состояние
                self.current_iterator = self.stack.pop()
                return next(self)
            else:
                # Выходить выше некуда, значит дошли до конца внешнего списка,
                # просто выбрасываем исключение (StopIteration) заново
                raise

        if type(item) is not list:
            return item
        else:
            # Запоминаем текущее состояние
            self.stack.append(self.current_iterator)
            # Входим во вложенный список, он становится текущим списком
            self.current_iterator = iter(item)
            return next(self)


list_of_lists_2 = [
    [['a'], ['b', 'c']],
    ['d', 'e', [['f'], 'h'], False],
    [1, 2, None, [[[[['!']]]]], []],
]

print(list(FlatIterator(list_of_lists_2)))

Также, чтобы код был точно не рекурсивным, нужно избавиться от вызовов next(self) из метода __next__ (по сути метод вызывает сам себя). Для этого нужно обернуть код в методе в while True и вызовы next(self) заменить на continue:

class FlatIterator:
    ...

    def __next__(self):
        while True:
            try:
                item = next(self.current_iterator)
            except StopIteration:
                if self.stack:
                    # Дошли до конца вложенного списка,
                    # выходим из него во внешний, "вспоминаем" состояние
                    self.current_iterator = self.stack.pop()
                    continue
                else:
                    # Выходить выше некуда, значит дошли до конца внешнего списка,
                    # просто выбрасываем исключение (StopIteration) заново
                    raise

            if type(item) is not list:
                return item
            else:
                # Запоминаем текущее состояние
                self.stack.append(self.current_iterator)
                # Входим во вложенный список, он становится текущим списком
                self.current_iterator = iter(item)
                # continue

...
▲ 1

Мой ответ повторяет другой, уже принятый. Всё же я покажу два решения. Они довольно компактные.

Первое хранит стек из пар список/индекс. Индекс ссылается на первый не обработанный элемент в списке. При создании итератора в стек кладётся единственный элемент - корневой список.

Метод __next__ в цикле отыскивает следующий элемент-не список. Работа ведётся с вершиной стека. Если очередной элемент - список, в стек добавляется новая пара, иначе элемент возвращается наружу. Если список на вершине стека исчерпан, он снимается со стека. Если стек опустел, бросается исключение:

class FlatIterator:
    def __init__(self, list_):
        self._stack = [[list_, 0]]

    def __iter__(self):
        return self

    def __next__(self):
        while self._stack:
            list_, i = self._stack[-1]
            if i < len(list_):
                self._stack[-1][1] += 1
                if not isinstance(list_[i], list):
                    return list_[i]
                self._stack.append([list_[i], 0])
            else:
                self._stack.pop()
        raise StopIteration

Если вы понимаете как работает вариант со списками и индексами, вы разберётесь как работает стек итераторов. Это более питонический метод, но смысл тот же:

class FlatIterator:
    def __init__(self, list_):
        self._stack = [iter(list_)]

    def __iter__(self):
        return self

    def __next__(self):
        while self._stack:
            try:
                item = next(self._stack[-1])
            except StopIteration:
                self._stack.pop()
                continue
            if not isinstance(item, list):
                return item
            self._stack.append(iter(item))
        raise StopIteration

P.S. Оба кода достаточно сложны. И это причина почему так не пишут - рекурсивный генератор проще написать и проще понять. Но есть одна ситуация когда без стека не обойтись - если уровень вложенности списков - 1000 или больше, то переполнится стек интерпретатора. Так что версия со стеком имеет право на существование.