DSA Flattening a Binary Tree

The problem

Flattening a binary tree rewires the tree in place into one chain that follows preorder. Preorder is a walk that visits a node before either of its children. In place means the same nodes are reused and no second tree is built.

In the finished chain, each right pointer holds the next node in preorder. Each left pointer holds None. The values themselves never change.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def preorder_values(node, order):
    if node is None:
        return
    order.append(node.value)
    preorder_values(node.left, order)
    preorder_values(node.right, order)

def read_chain(root):
    values = []
    node = root
    while node is not None:
        values.append(node.value)
        node = node.right
    return values

tree = TreeNode(5,
                TreeNode(3, TreeNode(1), TreeNode(4)),
                TreeNode(8, TreeNode(7), TreeNode(9)))

small = TreeNode(2, TreeNode(1), TreeNode(3))
for name, root, expected in [("main tree", tree, [5, 3, 1, 4, 8, 7, 9]),
                             ("small tree", small, [2, 1, 3]),
                             ("single node", TreeNode(9), [9]),
                             ("empty tree", None, [])]:
    order = []
    preorder_values(root, order)
    print(name, "- preorder:", order, " chain after flattening:", expected)

The main tree is the worked case from the source lesson. Its chain must read 5, 3, 1, 4, 8, 7, 9. Any other order is a wrong flattening, however tidy the chain looks.

Intuition: the chain order is the preorder list

The first idea reads the tree in preorder and writes the values down. That list is the exact order the chain must keep. Flattening is then a rewiring job, not a sorting job.

Walking right pointers alone is the trap. The root's right child is 8, so a right-only walk jumps from 5 to 8 at once. The whole left subtree would drop out of the chain.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def preorder_values(node, order):
    if node is None:
        return
    order.append(node.value)
    preorder_values(node.left, order)
    preorder_values(node.right, order)

def read_chain(root):
    values = []
    node = root
    while node is not None:
        values.append(node.value)
        node = node.right
    return values

tree = TreeNode(5,
                TreeNode(3, TreeNode(1), TreeNode(4)),
                TreeNode(8, TreeNode(7), TreeNode(9)))

order = []
preorder_values(tree, order)
print("preorder:", order)
print("root right child now:", tree.right.value)
print("values a right-only walk would reach first:", tree.value, "then", tree.right.value)
print("left subtree values at risk:", order[1:4])

The values at risk are 3, 1 and 4. A correct method moves each left subtree into the chain before the old right subtree. That move is the heart of the problem.

Brute force: collect the nodes, then rewire them

The brute force collects the node objects in preorder into a list. It then walks the list once. Each node's left pointer becomes None, and each right pointer aims at the next node in the list.

The method is simple and it works. Its cost is the list itself. The list holds every node in the tree before any rewiring starts.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def preorder_nodes(node, nodes):
    if node is None:
        return
    nodes.append(node)
    preorder_nodes(node.left, nodes)
    preorder_nodes(node.right, nodes)

def read_chain(root):
    values = []
    lefts = []
    node = root
    while node is not None:
        values.append(node.value)
        lefts.append(node.left)
        node = node.right
    return values, lefts

tree = TreeNode(5,
                TreeNode(3, TreeNode(1), TreeNode(4)),
                TreeNode(8, TreeNode(7), TreeNode(9)))

nodes = []
preorder_nodes(tree, nodes)
print("nodes stored in the list:", len(nodes))
for index, node in enumerate(nodes):
    node.left = None
    node.right = nodes[index + 1] if index + 1 < len(nodes) else None
values, lefts = read_chain(tree)
print("chain:", values)
print("left pointers that are None:", sum(1 for left in lefts if left is None))

Seven nodes need a seven-entry list. The chain afterwards matches preorder exactly. The list, not the rewiring, is what a better method has to remove.

The optimal solution: splice each left chain in and return the tail

Each node flattens both of its subtrees, then splices its flattened left chain between itself and its old right child, and returns the last node of its flattened subtree. That one sentence is the whole method.

The returned node is called the tail. The parent needs the tail to attach the next chain segment. Returning it keeps every splice a fixed amount of pointer work.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def flatten(node):
    if node is None:
        return None
    left_tail = flatten(node.left)
    right_tail = flatten(node.right)
    old_right = node.right
    if node.left is not None:
        node.right = node.left
        left_tail.right = old_right
        node.left = None
    if right_tail is not None:
        return right_tail
    if left_tail is not None:
        return left_tail
    return node

def preorder_values(node, order):
    if node is None:
        return
    order.append(node.value)
    preorder_values(node.left, order)
    preorder_values(node.right, order)

def read_chain(root):
    values = []
    node = root
    while node is not None:
        values.append(node.value)
        node = node.right
    return values

def fresh_main():
    return TreeNode(5,
                    TreeNode(3, TreeNode(1), TreeNode(4)),
                    TreeNode(8, TreeNode(7), TreeNode(9)))

for name, root, expected in [("main tree", fresh_main(), [5, 3, 1, 4, 8, 7, 9]),
                             ("small tree", TreeNode(2, TreeNode(1), TreeNode(3)), [2, 1, 3]),
                             ("single node", TreeNode(9), [9]),
                             ("empty tree", None, [])]:
    flatten(root)
    chain = read_chain(root)
    print(name, "-> chain:", chain, "expected:", expected, "match:", chain == expected)

All four cases from the start of the article pass. No list of nodes is built anywhere. The main tree flattens to the same chain the list method produced.

Walkthrough: splices in the main tree, bottom up

The walkthrough flattens the main tree with a traced copy of the splice method. A line prints at each node after its subtrees are flat. The line names the splice that node performs, or reports that it has none.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def flatten_traced(node):
    if node is None:
        return None
    left_tail = flatten_traced(node.left)
    right_tail = flatten_traced(node.right)
    old_right = node.right
    if node.left is not None:
        print("node", node.value, "- left chain ends at", left_tail.value,
              "- link it to old right child",
              old_right.value if old_right is not None else None)
        node.right = node.left
        left_tail.right = old_right
        node.left = None
    else:
        print("node", node.value, "- no left child, nothing to splice")
    if right_tail is not None:
        return right_tail
    if left_tail is not None:
        return left_tail
    return node

def read_chain(root):
    values = []
    node = root
    while node is not None:
        values.append(node.value)
        node = node.right
    return values

tree = TreeNode(5,
                TreeNode(3, TreeNode(1), TreeNode(4)),
                TreeNode(8, TreeNode(7), TreeNode(9)))

flatten_traced(tree)
print("chain after:", read_chain(tree))

Node 3 splices node 1 in before node 4. Node 8 splices node 7 in before node 9. Node 5 performs the last splice, linking node 4, the end of its flattened left subtree, to its old right child 8.

Read the chain from the root's right pointers at the end. It matches the preorder list from the first section exactly. Every left pointer along the way is None.

Complexity

Both methods visit each node a fixed number of times, so both take O(n) time for n nodes. The list method stores every node before rewiring, so its extra space is O(n).

The splice method stores no list, so its extra space is its recursion stack. The stack holds one root-to-leaf path at a time, so its space is O(h) for height h.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def flatten(node):
    if node is None:
        return None
    left_tail = flatten(node.left)
    right_tail = flatten(node.right)
    old_right = node.right
    if node.left is not None:
        node.right = node.left
        left_tail.right = old_right
        node.left = None
    if right_tail is not None:
        return right_tail
    if left_tail is not None:
        return left_tail
    return node

def preorder_nodes(node, nodes):
    if node is None:
        return
    nodes.append(node)
    preorder_nodes(node.left, nodes)
    preorder_nodes(node.right, nodes)

def flatten_measured(node, depth, max_depth):
    if node is None:
        return None
    if depth > max_depth[0]:
        max_depth[0] = depth
    left_tail = flatten_measured(node.left, depth + 1, max_depth)
    right_tail = flatten_measured(node.right, depth + 1, max_depth)
    old_right = node.right
    if node.left is not None:
        node.right = node.left
        left_tail.right = old_right
        node.left = None
    if right_tail is not None:
        return right_tail
    if left_tail is not None:
        return left_tail
    return node

def complete_tree(values):
    if not values:
        return None
    mid = len(values) // 2
    return TreeNode(values[mid], complete_tree(values[:mid]),
                    complete_tree(values[mid + 1:]))

for n in [7, 15, 31]:
    tree = complete_tree(list(range(1, n + 1)))
    nodes = []
    preorder_nodes(tree, nodes)
    max_depth = [0]
    flatten_measured(tree, 1, max_depth)
    print("nodes", n, "- brute list entries:", len(nodes),
          " optimal max stack depth:", max_depth[0])

The counts prove the space claim. The list grows with the node count itself. The stack depth grows only when the tree grows a level, which takes a doubling of the nodes here.

Edge cases and pitfalls

An empty tree flattens to an empty chain. A single node flattens to itself, with both pointers holding None. A tree that is already a right chain comes through unchanged.

Forgetting to clear the left pointers is the classic mistake. A chain read through right pointers can look right while left pointers still point into it. Any later walk that follows left pointers then visits nodes twice.

class TreeNode:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def flatten(node):
    if node is None:
        return None
    left_tail = flatten(node.left)
    right_tail = flatten(node.right)
    old_right = node.right
    if node.left is not None:
        node.right = node.left
        left_tail.right = old_right
        node.left = None
    if right_tail is not None:
        return right_tail
    if left_tail is not None:
        return left_tail
    return node

def preorder_values(node, order):
    if node is None:
        return
    order.append(node.value)
    preorder_values(node.left, order)
    preorder_values(node.right, order)

def read_chain(root):
    values = []
    node = root
    while node is not None:
        values.append(node.value)
        node = node.right
    return values

def left_count(root):
    count = 0
    node = root
    while node is not None:
        if node.left is not None:
            count += 1
        node = node.right
    return count

single = TreeNode(9)
flatten(single)
print("single node chain:", read_chain(single), "last right:", single.right)
left_chain = TreeNode(1, TreeNode(2, TreeNode(3)))
flatten(left_chain)
print("left chain flattened:", read_chain(left_chain))
right_chain = TreeNode(1, None, TreeNode(2, None, TreeNode(3)))
flatten(right_chain)
print("right chain flattened:", read_chain(right_chain))
print("empty tree chain:", read_chain(None))
main = TreeNode(5, TreeNode(3, TreeNode(1), TreeNode(4)), TreeNode(8))
flatten(main)
print("non-None left pointers after flattening main:", left_count(main))

The last count is the check that matters. Zero non-None left pointers remain after a correct flattening. The last node's right pointer holds None, because preorder gives it no next node.