Problem overview
The method flatten receives the root of a binary tree and rewires it in place into a "linked list" built from the same TreeNode objects. Every node's right pointer must point to the next node of the list, every left pointer must be nil, and the list must follow the preorder of the original tree: a node, then its left subtree, then its right subtree. The method returns nothing, so the change has to happen to the nodes themselves.
In the example, the tree [1, 2, 5, 3, 4, null, 6] has 1 at the root, children 2 and 5, and 2 has children 3 and 4 while 5 has the child 6. Its preorder is 1, 2, 3, 4, 5, 6, so the flattened tree is a chain going through right pointers in that order. An empty tree stays empty, and the single node [0] stays as it is.
Approach: build the list from the back
A natural first attempt flattens the left and right subtrees, then splices the right list onto the end of the left one. That works but needs a walk to find the end of the left list. There is a neater trick: process the nodes in the reverse of preorder, so the last node of the list is handled first. Then each node only has to point at the node processed just before it, which is the head of the list built so far.
Reverse preorder is right subtree, then left subtree, then the node. A recursive function that visits in that order, and remembers the most recently processed node in a field, does the whole job. For each node, lastNode is the first node of what comes after it in preorder, so setting root.right = lastNode links it in the right place. Clearing left and recording the node as the new lastNode finishes the step.
Tracing the example: the first node finished is 6, the last in preorder. Its right becomes nil (nothing was processed before it) and it becomes lastNode. Next comes 5, whose right subtree is already done: it points to 6. Then 4 points to 5, 3 points to 4, and 2 points to 3. Last, 1 points to 2. The chain is 1, 2, 3, 4, 5, 6.
Step by step
- Declare
lastNodeas aTreeNode?property ofSolution. Every recursive call has to see and update the same value, which a property provides. - Reset it in
flatten. ASolutionobject may be used for more than one tree, and the property keeps its value between calls, so each call starts withlastNode = nil. - Call
dfs(root)to do the work. - Declare
dfswith aTreeNode?parameter. It returns nothing and changes the nodes in place. - Stop on
nil. An empty subtree has nothing to flatten. - Flatten the right subtree first:
dfs(root!.right). The!is safe after the check. - Then flatten the left subtree:
dfs(root!.left). AfterwardslastNodeis the head of the left list, which links onto the right list built before it. - Point
root!.rightatlastNode. The current node now sits in front of everything flattened so far. - Set
root!.lefttonil. Nothing is lost: the left subtree is now reachable through therightpointers. - Record
rootaslastNode, so the node processed next links to it.
The solution in Swift
class Solution {
func dfs(_ root: TreeNode?) {
if root == nil {
return
}
dfs(root!.right)
dfs(root!.left)
root!.right = lastNode
root!.left = nil
lastNode = root
}
var lastNode: TreeNode?
func flatten(_ root: TreeNode?) {
lastNode = nil
dfs(root)
}
}
Swift details worth knowing:
- Properties can be declared after the methods that use them. A class's members are visible to each other in any order, so
dfscan uselastNodeeven though it is declared lower down. This is not true of top-level code outside a type. - Default value. A
varproperty of optional type starts asnilwithout an initializer, so the class needs noinit. The reset inflattenis still needed for the second and later calls. - Classes are references.
TreeNodeis a class, soroot!.right = lastNodechanges the real node. Assigning to the properties of a parameter is allowed, because only the reference is constant, not the object it points to. - Force unwrapping.
root!.left = nilwould crash ifrootwerenil; the early return guarantees it is not.
Complexity
Time: O(n). Each node is processed once with constant work, and each nil child adds one call that returns at once.
Space: O(h), where h is the height of the tree, for the call stack. No extra list is built; the nodes are rewired where they are. A tree shaped like a chain has h = n.
Edge cases and common mistakes
- Empty tree.
flatten(nil)returns immediately and nothing changes. - Single node. Its
rightbecomesnil(the resetlastNode) andleftis alreadynil. - Visiting left before right. The list would come out in the wrong order, since each node would link to a node that is not its preorder successor.
- Forgetting to clear
left. The result would still have left pointers, and the tree would not be a list. - Forgetting the reset. With a second tree, the last node of the new list would point into the old one.
- Linking before recursing. Changing
root!.rightbefore the right subtree is flattened would lose that subtree.
Testing
[1, 2, 5, 3, 4, null, 6]becomes the chain 1, 2, 3, 4, 5, 6.[1, null, 2, 3]becomes the chain 1, 2, 3.[0]stays the single node 0.- An empty tree stays empty.
