Problem overview
The function flatten receives root, a mutable reference to an Option<Box<TreeNode>>, and rearranges the tree in place into a "linked list" made of the same TreeNode type. In the result every node's left is None, its right points to the next node of the list, and the nodes are in the same order as a preorder traversal of the original tree: node, then left subtree, then right subtree.
For the tree [1, 2, 5, 3, 4, null, 6] the preorder is 1, 2, 3, 4, 5, 6. The flattened tree is [1, null, 2, null, 3, null, 4, null, 5, null, 6]: a chain going down the right side. An empty tree stays empty, and the tree [0] stays [0].
Approach: build the list backwards
A first idea is to collect the nodes in preorder into a vector and then relink them one after another. That works, but it needs extra space for the vector and a way to take the nodes out of the tree.
A neater idea uses the order of work. Imagine you already have the flattened list for everything that comes after some node in preorder. Then to flatten that node's part of the tree, you only need to put the node in front of that list. Preorder is node, left, right, so what comes after a node is its left subtree, then its right subtree, then whatever followed the whole subtree.
So the recursion processes the tree in the opposite order: right subtree first, then left subtree, then the node. The variable last_node always holds the head of the list built so far. After the right subtree has been handled, it holds the right subtree's list. After the left subtree has been handled, the left subtree's nodes sit in front of it, so it holds the left list followed by the right list. Finally the node itself takes that list as its right and becomes the new head.
On [1, 2, 5, 3, 4, null, 6] the nodes are linked in the order 6, 5, 4, 3, 2, 1, each one placed in front of the list: after 6 the list is 6, after 5 it is 5, 6, after 4 it is 4, 5, 6, and so on until 1, 2, 3, 4, 5, 6.
Step by step
- Create
last_nodeinflattenas a mutable variable starting atNone. It holds the list flattened so far, and nothing has been flattened yet. TheSolutionstruct has no fields to keep it in, so the recursion receives it as a mutable reference. - Call the helper
Self::dfs(root.take(), &mut last_node).take()moves the whole tree out of theOption, leavingNonebehind, sodfscan own the nodes and relink them. - Put the result back. When
dfsreturns,last_nodeholds the flattened tree. Assigning it to*roothands it back to the caller. - Declare
dfswithroot: Option<Box<TreeNode>>, which it owns, andlast_node: &mut Option<Box<TreeNode>>. It returns nothing. - Handle the empty case.
if let Some(mut node) = rootskips everything when the subtree is empty. Themutlets the code change the node's children. - Flatten the right subtree first, with
node.right.take(), passinglast_nodeon as it is. - Flatten the left subtree the same way with
node.left.take(). Taking the left child also leavesnode.leftasNone, which every node of the result needs. - Link the list after the node. Set
node.righttolast_node.take(). - Make the node the head. Store
Some(node)in*last_node.
The solution in Rust
impl Solution {
fn dfs(root: Option<Box<TreeNode>>, last_node: &mut Option<Box<TreeNode>>) {
if let Some(mut node) = root {
Self::dfs(node.right.take(), last_node);
Self::dfs(node.left.take(), last_node);
node.right = last_node.take();
*last_node = Some(node);
}
}
fn flatten(root: &mut Option<Box<TreeNode>>) {
let mut last_node = None;
Self::dfs(root.take(), &mut last_node);
*root = last_node;
}
}
Rust details worth noticing:
Option::take(). It replaces the value in anOptionwithNoneand returns the old value. That is how code moves a child out of a node that is only borrowed, because assigningnode.rightaway without leaving something behind isn't allowed.- Ownership flows down.
dfstakesrootby value, so each call owns its subtree. A node is moved into*last_nodein the last line, and the next call moves it back out withtake(). Nothing is copied, only the boxes are moved. &mutpassed along.last_nodeis already a mutable reference, so the recursive calls pass it directly. Passing it as it is lends it for the call, and it is usable again afterwards. Writing&mut last_nodethere is an error: it tries to borrow the parameter itself mutably, and the parameter isn't declaredmut.*last_nodeand*root. The star reaches theOptionthat the reference points to. Assigning to*rootchanges the caller's own variable, which is how a function with no return value delivers its result.mut nodein the pattern. The binding is mutable so thatnode.right = ...is allowed. Withoutmutthe compiler rejects the assignment.- Associated functions. Neither function takes
self, so the recursive calls are writtenSelf::dfs.
Complexity
Time: O(n). Each node is visited once, and each visit does a constant number of take and assignment operations.
Space: O(h), where h is the height of the tree, for the recursion stack. For a balanced tree that is about log n, and for a tree that leans entirely to one side it is n. No other structure is allocated, because the nodes are relinked in place.
Edge cases and common mistakes
- An empty tree.
root.take()givesNone,dfsdoes nothing, and*rootbecomesNoneagain. - A single node. It is flattened into itself:
rightbecomesNoneand the node is stored inlast_node. - A node with only a left child. For
[1, 2]the list is 1, 2 on the right side: the left child moves toright, and the oldleftbecomesNone. - A node with only a right child. For
[1, null, 2]the tree is already flat. The right subtree is flattened first, the left one is empty, and the node links to it. - Visiting the left subtree before the right. This builds the list in the wrong order, because the list is built from its last node to its first.
- Forgetting to clear the left link. Using
node.left.take()clears it. Passing a clone, or reading the child without taking it, leaves old left children in the result. - Very deep trees. The recursion goes as deep as the tree is tall, so a tree shaped like a very long chain can overflow the stack.
Testing
[1, 2, 5, 3, 4, null, 6]gives[1, null, 2, null, 3, null, 4, null, 5, null, 6].[]gives[].[0]gives[0].[1, 2, null, 3]gives[1, null, 2, null, 3].[1, null, 2, 3]gives[1, null, 2, null, 3]: here 3 is the left child of 2.
