Flatten a Binary Tree to a Linked List

Given the root of a binary tree, flatten it in place into a "linked list": every node's left is null, and its right is the next node in pre-order (node, then left subtree, then right subtree).

LeetCode: 114. Flatten Binary Tree to Linked List

Example

Text

      1
     / \
    2   5          ->   1 -> 2 -> 3 -> 4 -> 5 -> 6   (each arrow is a right link)
   / \   \
  3   4   6

Brute force: rebuild from the pre-order values

Record the values in pre-order, then build a new right-leaning chain from them.

A tree node is an object with an int val and two reference fields, left and right. A missing child is null, and the empty tree is a null root. The fields are mutable, so the code rewires nodes by assigning to them.

Java

import java.util.ArrayList;
import java.util.List;

public class Main {
    static class TreeNode {
        int val;
        TreeNode left, right;
        TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }

    static TreeNode node(int val, TreeNode left, TreeNode right) {
        return new TreeNode(val, left, right);
    }

    static TreeNode leaf(int val) {
        return new TreeNode(val, null, null);
    }

    static void preorder(TreeNode root, List<Integer> out) {
        if (root == null) return;
        out.add(root.val);
        preorder(root.left, out);
        preorder(root.right, out);
    }

    // Returns the root of a new chain; the old nodes are left as they were.
    static TreeNode flatten(TreeNode root) {
        List<Integer> values = new ArrayList<>();
        preorder(root, values);
        TreeNode next = null;
        for (int i = values.size() - 1; i >= 0; i--) {
            next = node(values.get(i), null, next);
        }
        return next;
    }

    // Follows right links; throws if a left link is left over.
    static List<Integer> chain(TreeNode root) {
        List<Integer> values = new ArrayList<>();
        for (TreeNode cur = root; cur != null; cur = cur.right) {
            if (cur.left != null) throw new AssertionError("left link left over at " + cur.val);
            values.add(cur.val);
        }
        return values;
    }

    public static void main(String[] args) {
        TreeNode root = node(1, node(2, leaf(3), leaf(4)), node(5, null, leaf(6)));
        System.out.println(chain(flatten(root)));
    }
}

Output

[1, 2, 3, 4, 5, 6]

Time: O(n). Space: O(n), and it allocates new nodes, so it is not in place.

Optimal: rewire each left subtree

Walk down the right links from the root. At each node n that has a left subtree:

Find the rightmost node of the left subtree. In pre-order, it is the last node visited before n's right subtree.
Attach n's right subtree to that rightmost node.
Move the left subtree to n.right, leaving n.left as null.

Then continue with n.right. Nodes further down are handled when the walk reaches them.

Text

    1                1                  1
   / \                \                  \
  2   5      ->        2         ->       2
 / \   \              / \                  \
3   4   6            3   4                  3 ...
                          \
                           5
                            \
                             6

Java

import java.util.ArrayList;
import java.util.List;

public class Main {
    static class TreeNode {
        int val;
        TreeNode left, right;
        TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }

    static TreeNode node(int val, TreeNode left, TreeNode right) {
        return new TreeNode(val, left, right);
    }

    static TreeNode leaf(int val) {
        return new TreeNode(val, null, null);
    }

    static void flatten(TreeNode root) {
        TreeNode cur = root;
        while (cur != null) {
            if (cur.left != null) {
                TreeNode rightmost = cur.left;
                while (rightmost.right != null) {
                    rightmost = rightmost.right;
                }
                rightmost.right = cur.right; // the old right subtree follows the left subtree
                cur.right = cur.left;        // the left subtree becomes the right link
                cur.left = null;
            }
            cur = cur.right;
        }
    }

    // Follows right links; throws if a left link is left over.
    static List<Integer> chain(TreeNode root) {
        List<Integer> values = new ArrayList<>();
        for (TreeNode cur = root; cur != null; cur = cur.right) {
            if (cur.left != null) throw new AssertionError("left link left over at " + cur.val);
            values.add(cur.val);
        }
        return values;
    }

    public static void main(String[] args) {
        TreeNode root = node(1, node(2, leaf(3), leaf(4)), node(5, null, leaf(6)));
        flatten(root);
        System.out.println(chain(root));

        TreeNode leftOnly = node(1, node(2, node(3, null, null), null), null);
        flatten(leftOnly);
        System.out.println(chain(leftOnly));

        TreeNode single = leaf(7);
        flatten(single);
        flatten(null);
        System.out.println(chain(single) + " " + chain(null));
    }
}

Output

[1, 2, 3, 4, 5, 6]
[1, 2, 3]
[7] []

Time: O(n): each node is on the path to a rightmost node at most once. Space: O(1).

Every step moves a reference from one field to another, so no node is copied or reallocated, and the structure stays a tree the whole time.

Edge cases

Empty tree or single node: unchanged.
Already a right chain: no node has a left subtree, so nothing moves.