430. Flatten a Multilevel Doubly Linked List

[mengyushi]

"""
# Definition for a Node.
class Node:
    def __init__(self, val, prev, next, child):
        self.val = val
        self.prev = prev
        self.next = next
        self.child = child
"""
class Solution:
    def flatten(self, head: 'Node') -> 'Node':

        '''
        Using Stack

        stack           processed

        [1]             None
        [2]             1
        [3]             1->2
        [4, 7]          1->2->3
        [4, 8]          1->2->3->7
        [4, 9, 11]      1->2->3->7->8
        [4, 9, 12]      1->2->3->7->8->11
        [4, 9]          1->2->3->7->8->11->12
        [4, 10]         1->2->3->7->8->11->12->9
        [4]             1->2->3->7->8->11->12->9->10
        [5]             1->2->3->7->8->11->12->9->10->4
        [6]             1->2->3->7->8->11->12->9->10->4->5
        []              1->2->3->7->8->11->12->9->10->4->5->6

        Time Complexity: O(N)
        Space Complexit: O(N)

        '''

        if not head:
            return None

        dummy, stack = Node(0, None, head, None), [head]
        prev = dummy

        while stack:
            node = stack.pop()
            node.prev, prev.next = prev, node

            if node.next:
                stack.append(node.next)
                node.next = None

            if node.child:
                stack.append(node.child)
                node.child = None

            prev = node

        dummy.next.prev = None
        return dummy.next