python помогите с обходом дерева в цикле

Нужно обойти дерево с тремя слоями, но это не бинарное дерево, в слое может быть много элементов и не у каждого может быть снизу элемент. Обойти нужно без ООП, а вложенными циклами внутри одной функции.

для 2 вывод должен быть [[3,4],[5,6,7,8,],[9,10,1]]
для 3 вывод должен быть [[5,6],[9],[]]