Problem overview
The function flattenTree receives the root of a binary tree, a *TreeNode, and rearranges that same tree in place. It returns nothing. Afterwards every node's Left pointer must be nil, and the Right pointers must chain the nodes into a list in preorder: the node, then its left subtree, then its right subtree.
In the first example the tree is [1,2,5,3,4,null,6]: 1 has children 2 and 5, 2 has children 3 and 4, and 5 has a right child 6. Preorder is 1, 2, 3, 4, 5, 6. After flattening, the root is still 1, 1's right child is 2, 2's right child is 3, and so on down to 6. An empty tree stays empty, and a single node stays as it is.
Approach: build the list backwards
A direct approach is to collect the nodes in preorder in a slice, then walk the slice and relink each node to the next. That works, but it needs extra memory proportional to the tree.
This lesson avoids that by working in reverse. The flattened list is a chain of Right pointers, and a chain is easiest to build from its tail: if you already hold the node that must come next, linking the current node in front of it takes one assignment. Preorder is node, left, right, so its reverse is right, left, node. A recursive search that visits the right subtree, then the left subtree, then the node itself, sees the nodes in exactly the reverse of preorder.
A single variable, lastFlattened, remembers the node processed most recently, which is the head of the list built so far. When the search reaches a node, both of its subtrees are done, and lastFlattened is the node that must follow it in preorder: the first node of the left subtree, or if there is none, the first node of the right subtree, or if there is none either, whatever came after the whole subtree. The node links to it and becomes the new head.
Step by step
- Declare a package-level variable
lastFlattenedof type*TreeNode, aboveflattenTree. It starts asnil. Go has no classes to hold shared state, so a package-level variable plays the role of a field. - Reset it at the start of
flattenTree. A package-level variable keeps its value between calls, so without the reset a second call would link its last node to a node of the previous tree. - Call
flattenDfs(root)to flatten the whole tree. - Declare
flattenDfswith one parameterroot *TreeNodeand no result. - Return on
nil. An empty subtree has nothing to flatten. - Flatten the right subtree first. It comes last in the list, and the list is built from the end.
- Flatten the left subtree. When this call returns,
lastFlattenedis the first node of the left part, which must come right afterroot. - Set
root.RighttolastFlattened. This links the node to what follows it. - Set
root.Lefttonil. The left subtree is now reachable through theRightchain, so nothing is lost. - Set
lastFlattenedtoroot. This node is the head of the part that is done, and the next node to be processed links to it.
The solution in Go
var lastFlattened *TreeNode
func flattenTree(root *TreeNode) {
lastFlattened = nil
flattenDfs(root)
}
func flattenDfs(root *TreeNode) {
if root == nil {
return
}
flattenDfs(root.Right)
flattenDfs(root.Left)
root.Right = lastFlattened
root.Left = nil
lastFlattened = root
}
Tracing the first example shows the order. The search goes right from 1 to 5 and on to 6. Node 6 has no children, so it points to nil and becomes the head. Node 5 then points to 6. Back at 2, the right subtree (4) is flattened first and then the left (3). Node 4 points to 5, node 3 points to 4, node 2 points to 3, and finally node 1 points to 2.
The Go details worth a closer look:
- Package-level state. A
varoutside any function lives for the whole program and starts at its zero value, which for a pointer isnil. Every function in the package can read and assign it. The price is that the function isn't re-entrant, which is whyflattenTreeresets it. - In-place work through pointers.
rootis a pointer, soroot.Right = ...changes the caller's node. Go writes the field access asroot.Rightfor pointers too, with no->operator. - Read before write. The recursive calls run before the assignments, so
root.Leftandroot.Rightstill point to the original children when they are used. Onceroot.Rightis overwritten, the right subtree is only reachable through the list, which is fine because it has already been flattened. - Calling a function defined later.
flattenTreecallsflattenDfsbefore the file defines it. Go allows that for package-level functions.
Complexity
Time: O(n) for n nodes. Each node is visited once, and each visit does a few pointer assignments. Missing children add one cheap call each.
Space: O(h) for the recursion, where h is the height of the tree. There is no extra list or slice: the nodes are relinked where they are. For a tree that is one long chain, h is n.
Edge cases and common mistakes
- An empty tree. With
root == nilthe helper returns immediately and nothing changes. - A single node. Both subtrees are empty, so
Rightis set tonil,Lefttonil, and the node is unchanged. - Visiting left before right. The list would come out in the wrong order: the nodes of the right subtree would be placed before those of the left subtree.
- Skipping the reset. Flattening a second tree after the first would leave that tree's last node pointing to the head of the first list.
- Forgetting
root.Left = nil. The list would still be correct along theRightpointers, but the left pointers would keep the old subtrees attached, so the structure still looks like a tree. - Returning the new root. The function returns nothing, and the root stays the same node. A tree where the root has a left child still starts at the same root after flattening.
Testing
- The tree
[1,2,5,3,4,null,6]becomes the chain 1, 2, 3, 4, 5, 6 along theRightpointers. - An empty tree stays empty.
- A single node with value 0 stays
[0]. - A root 1 with only a left child 2 becomes 1, then 2 as the right child, and the left pointer of 1 is
nil.