Start learning free

LESSON 336 OF 414 · BINARY TREES

Flatten Binary Tree in Scala: Reverse Preorder with a Field

Practice O(n) Read the full explanation

Given the root of a binary tree, flatten the tree into a "linked list":

  • The "linked list" should use the same tree node type (TreeNode), where the right child link points to the next node in the list and the left child link is always null.
  • The "linked list" should be in the same order as a pre-order traversal of the binary tree.

 

Example 1:

Input: root = [1,2,5,3,4,null,6]
Output: [1,null,2,null,3,null,4,null,5,null,6]

Example 2:

Input: root = []
Output: []

Example 3:

Input: root = [0]
Output: [0]
STEP 1 OF 10
    Ctrl + Enter to check

    SCALA · BINARY TREES

    Flatten Binary Tree in Scala, explained

    Problem overview

    The method flatten receives the root of a binary tree and rewires it in place into a "linked list" made of the same TreeNode objects. In the result, every node's right pointer leads to the next node of the list, and every left pointer is null. The list must follow the preorder of the original tree: a node, then its left subtree, then its right subtree.

    The method returns Unit, so it has no result to hand back. The caller's root is the head of the list afterwards.

    In the first example, the tree has root 1, children 2 and 5, node 2 has children 3 and 4, and node 5 has a right child 6. Its preorder is 1, 2, 3, 4, 5, 6, so after flattening, the root 1 leads through its right pointers to 2, 3, 4, 5 and 6, and no node has a left child. An empty tree stays empty, and a single node stays as it is.

    Approach: build the list from the back

    The natural plan is to link each node to its successor in preorder. The difficulty is that, when you stand at a node, its successor is the head of the left subtree's list, and the end of that list must then connect to the right subtree. Finding that end takes extra work.

    The lesson avoids the problem by visiting the nodes in the reverse of preorder: right subtree, then left subtree, then the node. In that order, the node visited just before the current one is always the node that must follow it in the list. So one field, lastNode, is enough. It holds the head of the list built so far.

    When the walk is at a node and has finished both subtrees, lastNode is the head of the list for the left subtree, which ends by leading into the right subtree's list, which in turn leads into everything that comes later in preorder. That is exactly what should follow the current node. So the node points its right at lastNode, clears its left, and becomes the new lastNode.

    On the first example, the nodes are processed in the order 6, 5, 4, 3, 2, 1. Node 6 points to null, 5 points to 6, 4 points to 5, 3 points to 4, 2 points to 3, and finally 1 points to 2. The list reads 1, 2, 3, 4, 5, 6, which matches the preorder.

    Step by step

    1. Declare lastNode as a var of type TreeNode in the class, starting at null. As a field, it is shared by every recursive call of dfs.
    2. Reset it in flatten. A field keeps its value between calls, so the same Solution object used for a second tree would otherwise start with a leftover node from the first.
    3. Call dfs(root), a helper that flattens a whole subtree. It is written next, and a method of a class can call another method wherever it is placed.
    4. Declare dfs with a TreeNode parameter and the return type Unit, since it rewires nodes instead of returning anything.
    5. Handle the base case. When root is null, return.
    6. Recurse on the right subtree first. After this call, the right side is flattened and lastNode is its head.
    7. Recurse on the left subtree. After this call, the left side is flattened and its list ends by leading into the right list, and lastNode is the head of the left list.
    8. Point root.right at lastNode. The node now links straight into the list built so far.
    9. Clear root.left. The left subtree is already reachable through right pointers, so nothing is lost.
    10. Record root in lastNode. The node is the new head of the list, and the node processed next will link to it.

    The solution in Scala

    class Solution {
    def dfs(root: TreeNode): Unit = {
        if (root == null) {
            return
        }
        dfs(root.right)
        dfs(root.left)
        root.right = lastNode
        root.left = null
        lastNode = root
    }

    var lastNode: TreeNode = null
    def flatten(root: TreeNode): Unit = {
        lastNode = null
        dfs(root)
    }
}

    The Scala details worth a closer look:

    • A var in the class body is a field. var lastNode: TreeNode = null belongs to each Solution object, and every method reads and writes the same variable. The order of members in a class doesn't matter, so dfs can use lastNode although it is declared below it.
    • The type is spelled out. Writing only = null would make Scala infer the type Null, which can't hold a node. The annotation : TreeNode gives the variable the type it needs.
    • The node's pointers are assignable. root.right = lastNode and root.left = null work because this lesson's TreeNode declares left and right as var fields.
    • Unit means in place. flatten returns nothing, so the work is visible only through the nodes the caller already holds.
    • return for the early exit. In a method returning Unit, a bare return leaves the method at once. It is the simplest way to write the null base case.

    Complexity

    Time: O(n) for n nodes. Each node is processed once, with a constant amount of work for its three assignments, plus one extra call for each missing child.

    Space: O(h) for the call stack, where h is the height of the tree, and O(1) for everything else. The nodes are reused and no new ones are created. A tree that is one long chain has height n, so the stack use can reach O(n) in the worst case.

    Edge cases and common mistakes

    • Empty tree. A null root makes dfs return at once, and flatten changes nothing.
    • Single node. The node's right is set to null (the reset value of lastNode), and its left stays null.
    • Only a left child. Root 7 with left child 8 becomes 7, then 8 on the right. The left pointer of 7 is cleared.
    • Left child that has a right child. Root 1 with right child 2 whose left child is 3 becomes 1, 2, 3, because 3 is the left subtree of 2 and comes before whatever follows.
    • Recursing left before right. That would link each node to the wrong successor and produce a different order than preorder.
    • Forgetting root.left = null. The list would still have the right order through right, but nodes would keep their old left children and the result wouldn't be a list.
    • Forgetting to reset lastNode. On a second call with the same Solution object, the last node of the new tree would point into the old list.

    Testing

    • Root 1 with left child 2 (children 3 and 4) and right child 5 (right child 6) becomes 1, 2, 3, 4, 5, 6 on the right pointers.
    • A null root stays empty.
    • The single node 0 stays 0.
    • Root 7 with left child 8 becomes 7, 8.
    • Root 1 with right child 2, which has a left child 3, becomes 1, 2, 3.
    Keep going with ScalaAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free