Flatten Multilevel Linked List

https://leetcode.com/problems/flatten-a-multilevel-doubly-linked-list/
https://www.geeksforgeeks.org/flatten-a-linked-list-with-next-and-child-pointers/

Problem Description

Imagine a linked list with not just a next node, but a down node as well. We want to write a function that would "Flatten" that linked list by taking all the downward segments and merging them up between their parent and their parent's next.

Sample inputs - Expected outputs

If we have something like this:

[1] -> [2] -> [3] -> [8] -> [10]
               |      |
	           |     [9]
			   |
			  [4] -> [5] -> [6]
			                 |
							[7]

after flatten, it should look like this:

[1] -> [2] -> [3] -> [4] -> [5] -> [6] -> [7] -> [8] -> [9] -> [10]

Solution

Perform Depth First Search using a Stack, where we always go deeper in down nodes, and then return to the next nodes.

Start the first node [1] as the iterating node:

• Start by putting the next node [2] and then the down node None of the first node [1] into the stack ([2, None]). As down node does not exist, we don't have to put it in the stack, so the stack will be ([2]).

While the stack is not empty:

• Remove the down pointer for this node [1]->down=None.
• Pop a node [2] from the stack ([2] -> []).
• Point the iterating node [1] to the popped node [1]->[2].
• Start iterating the next node [2]
• Add the children nodes of the iterating node [2]
  • add its next node to the stack [3]
  • add its down node to the stack None (in this case, no need to add it)

For each node, we add its next node to the stack and then the down node.
The iterated node will point towards the last node in the stack.

Alternatively we can keep a stack of all next nodes when we encounter a down node.

Complexities for flatten_chain (where n is number of nodes):

• Time complexity: O(n)
• Space complexity: O(n)

from __future__ import annotations

class Node:
	def __init__(self, value: int, next: Node = None, down: Node = None):
		self.value = value
		self.next = next
		self.down = down
	def __str__(self):
		return str(self.value)

chain = Node(1,Node(2,Node(3,
	Node(8,Node(10),Node(9)),
	Node(4,Node(5,Node(6,down=Node(7))))
)))


def flatten_chain(chain: Node):
	node = chain
	next_stack = []

	def add_children(node: Node):
		if node.next:
			next_stack.append(node.next)
		if node.down:
			next_stack.append(node.down)

	add_children(node)
	while len(next_stack) > 0:
		node.down = None
		node.next = next_stack.pop()
		node = node.next
		add_children(node)

flatten_chain(chain)

I got rejected because during the interview I did not mention that I am using DFS; I also wrote while node.down or node.next or len(next_stack) > 0 and my code was harder to debug. The interviewer understood my solution well, however. The above solution contains a bit neater code.