Balance a BST

▶

Steps

setup

Initialize nodes list and start inorder traversal

Create an empty list 'nodes' to store nodes in sorted order. Begin inorder traversal from the root node (value 1).

💡 This step sets up the data structure to collect nodes in ascending order, which is key to balancing.

Line:

nodes = []

def inorder(node: Optional[TreeNode]):
    if not node:
        return

💡 The inorder traversal will visit nodes in ascending order, which is essential for balanced BST construction.

traverse

Traverse left subtree of node 1 (null)

Attempt to traverse the left child of node 1, which is null, so return immediately.

💡 This confirms the base case of recursion where null children are skipped.

Line:

if not node:
    return

💡 Leaves and null children mark the end of recursion branches.

fill_cells

Visit node 1 and add to nodes list

Add node 1 to the nodes list as part of inorder traversal after left subtree is processed.

💡 Nodes are collected in ascending order during inorder traversal.

Line:

nodes.append(node)

💡 Node 1 is the smallest node so far and is added first.

traverse

Traverse right subtree of node 1 (node 2)

Move to the right child of node 1, which is node 2, to continue inorder traversal.

💡 Inorder traversal visits left, node, then right, so now we explore node 2's subtree.

Line:

inorder(node.right)

💡 Traversal proceeds to next larger node in BST.

traverse

Traverse left subtree of node 2 (null)

Attempt to traverse left child of node 2, which is null, so return immediately.

💡 Null children mark recursion base cases.

Line:

if not node:
    return

💡 Leaves are boundaries for recursion.

fill_cells

Visit node 2 and add to nodes list

Add node 2 to the nodes list after finishing its left subtree traversal.

💡 Nodes are appended in ascending order during inorder traversal.

Line:

nodes.append(node)

💡 Node 2 is the next smallest node and is added after node 1.

traverse

Traverse right subtree of node 2 (node 3)

Move to the right child of node 2, which is node 3, continuing inorder traversal.

💡 Traversal proceeds to next larger node in BST.

Line:

inorder(node.right)

💡 The traversal will visit nodes 3 and 4 next.

traverse

Traverse left subtree of node 3 (null)

Attempt to traverse left child of node 3, which is null, so return immediately.

💡 Null children mark recursion base cases.

Line:

if not node:
    return

💡 Leaves are boundaries for recursion.

fill_cells

Visit node 3 and add to nodes list

Add node 3 to the nodes list after finishing its left subtree traversal.

💡 Nodes are appended in ascending order during inorder traversal.

Line:

nodes.append(node)

💡 Node 3 is the next smallest node and is added after nodes 1 and 2.

traverse

Traverse right subtree of node 3 (node 4)

Move to the right child of node 3, which is node 4, continuing inorder traversal.

💡 Traversal proceeds to the largest node in the BST.

Line:

inorder(node.right)

💡 Node 4 will be the last node added to the list.

traverse

Traverse left subtree of node 4 (null)

Attempt to traverse left child of node 4, which is null, so return immediately.

💡 Null children mark recursion base cases.

Line:

if not node:
    return

💡 Leaves are boundaries for recursion.

fill_cells

Visit node 4 and add to nodes list

Add node 4 to the nodes list after finishing its left subtree traversal.

💡 Nodes are appended in ascending order during inorder traversal.

Line:

nodes.append(node)

💡 Node 4 is the largest node and is added last.

setup

Initialize stack for iterative balanced BST construction

Create a stack with initial tuple (0, 3, None, True) representing the full nodes array and no parent.

💡 Stack will manage subarray indices to build balanced BST iteratively.

Line:

stack = [(0, len(nodes) - 1, None, True)]
root = None

💡 Stack-driven approach replaces recursion for balanced BST construction.

insert

Pop stack and select middle node 2 as root

Pop (0,3,None,True) from stack. Calculate mid=1, select node 2 as root. Set root to node 2.

💡 Choosing middle node balances the tree by splitting nodes evenly.

Line:

left, right, parent, is_left = stack.pop()
mid = (left + right) // 2
node = nodes[mid]
root = node

💡 Middle node 2 balances left and right subtrees.

expand

Push right subtree indices (2,3) and left subtree indices (0,0) to stack

Push (2,3,node2,False) for right subtree and (0,0,node2,True) for left subtree onto stack.

💡 Stack keeps track of subarrays to process for building subtrees.

Line:

stack.append((mid + 1, right, node, False))
stack.append((left, mid - 1, node, True))

💡 Subtrees will be built by processing these subarray ranges.

insert

Pop stack and select middle node 3 as right child of 2

Pop (2,3,node2,False). Calculate mid=2, select node 3. Set node 3 as right child of node 2.

💡 Middle node of right subtree subarray becomes right child.

Line:

left, right, parent, is_left = stack.pop()
mid = (left + right) // 2
node = nodes[mid]
parent.right = node

💡 Right subtree is balanced by choosing middle node 3.

expand

Push right subtree indices (3,3) and left subtree indices (2,1) for node 3

Push (3,3,node3,False) for right subtree and (2,1,node3,True) for left subtree onto stack.

💡 Continue building subtrees for node 3 using subarray indices.

Line:

stack.append((mid + 1, right, node, False))
stack.append((left, mid - 1, node, True))

💡 Subtrees of node 3 will be constructed next.

insert

Pop stack and select middle node 4 as right child of 3

Pop (3,3,node3,False). Calculate mid=3, select node 4. Set node 4 as right child of node 3.

💡 Leaf node 4 becomes right child of node 3.

Line:

left, right, parent, is_left = stack.pop()
mid = (left + right) // 2
node = nodes[mid]
parent.right = node

💡 Rightmost node is placed correctly as leaf.

prune

Pop stack with invalid left subtree indices (2,1) and skip

Pop (2,1,node3,True). Since left > right, skip processing this invalid subarray.

💡 Subarrays with left > right represent empty subtrees and are ignored.

Line:

if left > right:
    continue

💡 Algorithm correctly handles empty subtrees by ignoring invalid ranges.

reconstruct

Balanced BST construction complete, return root node 2

Stack is empty, construction finished. Return root node 2 representing balanced BST.

💡 Final balanced BST root is ready for use.

Line:

return root

💡 Balanced BST root is node 2 with left child 1 and right subtree rooted at 3 with right child 4.

from typing import Optional, List

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def balanceBST(root: Optional[TreeNode]) -> Optional[TreeNode]:
    nodes = []  # STEP 1: Initialize list to collect nodes

    def inorder(node: Optional[TreeNode]):
        if not node:  # STEP 2: Base case for recursion
            return
        inorder(node.left)  # STEP 4,5,8,11: Traverse left subtree
        nodes.append(node)  # STEP 3,6,9,12: Add node to list
        inorder(node.right)  # STEP 4,7,10: Traverse right subtree

    inorder(root)  # STEP 1-12: Collect nodes in sorted order

    if not nodes:
        return None

    stack = [(0, len(nodes) - 1, None, True)]  # STEP 13: Initialize stack
    root = None

    while stack:
        left, right, parent, is_left = stack.pop()  # STEP 14,16,18,19: Pop stack
        if left > right:  # STEP 19: Skip invalid ranges
            continue
        mid = (left + right) // 2
        node = nodes[mid]
        node.left = node.right = None
        if parent:
            if is_left:
                parent.left = node
            else:
                parent.right = node
        else:
            root = node
        stack.append((mid + 1, right, node, False))  # STEP 15,17: Push right subtree
        stack.append((left, mid - 1, node, True))  # STEP 15,17: Push left subtree

    return root  # STEP 20: Return balanced BST root

📊

Balance a BST - Watch the Algorithm Execute, Step by Step

Watching the algorithm step-by-step reveals how the skewed BST is transformed into a balanced one by reorganizing nodes based on inorder traversal.

Current node: 1

1234

DFS Stack

1 (entered)

1234

DFS Stack

1 (left_done)

Current node: 1

1234

DFS Stack

1 (node_added)

Result: [1]

Current node: 2

1234

DFS Stack

2 (entered)

Result: [1]

1234

DFS Stack

2 (left_done)

Result: [1]

Current node: 2

1234

DFS Stack

2 (node_added)

Result: [1, 2]

Current node: 3

1234

DFS Stack

3 (entered)

Result: [1, 2]

1234

DFS Stack

3 (left_done)

Result: [1, 2]

Current node: 3

1234

DFS Stack

3 (node_added)

Result: [1, 2, 3]

Current node: 4

1234

DFS Stack

4 (entered)

Result: [1, 2, 3]

1234

DFS Stack

4 (left_done)

Result: [1, 2, 3]

Current node: 4

1234

Result: [1, 2, 3, 4]

1234

Result: [1, 2, 3, 4]

Current node: 2

1234

Result: [1, 2, 3, 4]

Current node: 2

1234

Result: [1, 2, 3, 4]

Current node: 3

1234

Result: [1, 2, 3, 4]

Current node: 3

1234

Result: [1, 2, 3, 4]

Current node: 4

1234

Result: [1, 2, 3, 4]

1234

Result: [1, 2, 3, 4]

Current node: 2

1234

Result: [1, 2, 3, 4]

Return: 2

Key Takeaways

✓ Inorder traversal collects nodes in sorted order, which is the foundation for balancing the BST.

This insight is hard to see from code alone because the traversal and balancing are separate phases.

✓ Choosing the middle node of the sorted nodes array as root at each step ensures the tree is balanced.

Visualizing the middle selection clarifies how the tree height is minimized.

✓ Using a stack to iteratively build the balanced BST replaces recursion and manages subtree construction explicitly.

Seeing the stack operations helps understand how the algorithm controls subtree boundaries.

Practice

1. Consider the following BSTIterator implementation using Morris traversal:

class BSTIterator:
    def __init__(self, root):
        self.current = root
        self.next_val = None
        self._advance()

    def _advance(self):
        while self.current:
            if not self.current.left:
                self.next_val = self.current.val
                self.current = self.current.right
                return
            else:
                pre = self.current.left
                while pre.right and pre.right != self.current:
                    pre = pre.right
                if not pre.right:
                    pre.right = self.current
                    self.current = self.current.left
                else:
                    pre.right = None
                    self.next_val = self.current.val
                    self.current = self.current.right

Given the BST with nodes 2 (root), 1 (left child), and 3 (right child), what is the value of next_val after the first call to _advance() during initialization?

easy

A. 3

B. 2

C. 1

D. null

Solution

1. Step 1: Trace _advance() starting at root=2

   Since root has a left child (1), find predecessor of 2 in left subtree: node 1 has no right child, so set 1.right = 2 and move current to 1.

2. Step 2: At current=1, no left child, so set next_val=1 and move current to 1.right which points back to 2 (thread).

   Return from _advance() with next_val=1.

3. Final Answer:

   Option C -> Option C

4. Quick Check:

   First inorder element is 1, matches next_val after first advance [OK]

Hint: Morris traversal first yields leftmost node [OK]

Common Mistakes:

• Assuming next_val is root value initially
• Forgetting thread creation
• Confusing next_val with current node

2. You are given a binary tree where exactly two nodes have been swapped by mistake, violating the binary search tree property. Which approach guarantees an optimal solution to recover the tree without changing its structure or using extra space beyond a constant amount?

easy

A. Use a bottom-up dynamic programming approach to reconstruct the BST by comparing subtrees and swapping nodes.

B. Perform an inorder traversal, store values in a list, sort the list, then overwrite the tree nodes with sorted values.

C. Apply a greedy approach by swapping any two nodes that violate the BST property during a single inorder traversal pass.

D. Use Morris inorder traversal to detect the two swapped nodes during traversal and swap their values, achieving O(1) space complexity.

Solution

1. Step 1: Understand the problem constraints

   The problem requires recovering a BST where two nodes are swapped, without extra space and without changing the tree structure.

2. Step 2: Identify the approach that meets constraints

   Morris inorder traversal allows inorder traversal without recursion or stack, detecting swapped nodes on the fly and swapping their values, achieving O(1) space.

3. Final Answer:

   Option D -> Option D

4. Quick Check:

   Morris traversal is the only approach with O(1) space and no structural changes [OK]

Hint: Morris traversal enables O(1) space recovery [OK]

Common Mistakes:

• Assuming sorting values is optimal despite extra space
• Believing greedy single pass fixes all cases
• Confusing DP with tree recovery

3. You are given a binary tree where each node has a unique integer value. You need to find whether a given value exists in the tree efficiently by leveraging the tree's ordering properties. Which approach guarantees the fastest search in the worst case?

easy

A. Use the BST property to decide whether to search left or right subtree at each node, recursively or iteratively.

B. Perform a breadth-first search (BFS) to find the value level by level.

C. Traverse the entire tree recursively checking each node until the value is found.

D. Use dynamic programming to store previously searched subtrees to avoid repeated work.

Solution

1. Step 1: Understand the problem constraints

   The tree is a BST, so left subtree nodes are less than the node, right subtree nodes are greater.

2. Step 2: Identify the optimal search approach

   Using the BST property allows skipping half of the tree at each step, leading to O(h) time complexity.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   BST property guides search direction efficiently [OK]

Hint: BST property enables directed search, not full traversal [OK]

Common Mistakes:

• Thinking BFS or full traversal is optimal for BST search

4. Consider the following buggy code for inserting a value into a BST. Which line contains the subtle bug that causes the insertion to fail silently in some cases?

def insertIntoBST(root, val):
    if root is None:
        return TreeNode(val)
    if val < root.val:
        insertIntoBST(root.left, val)
    else:
        insertIntoBST(root.right, val)
    return root

medium

A. Line 2: if root is None

B. Line 4: if val < root.val

C. Line 7: insertIntoBST(root.right, val)

D. Line 5: insertIntoBST(root.left, val)

Solution

1. Step 1: Analyze recursive calls

   Recursive calls to insertIntoBST on root.left or root.right do not update the parent's child pointer.

2. Step 2: Identify missing assignment

   Lines 5 and 7 call insertIntoBST but do not assign the returned subtree back to root.left or root.right, so insertion is lost.

3. Final Answer:

   Option D -> Option D

4. Quick Check:

   Missing assignment causes silent insertion failure [OK]

Hint: Recursive insert must assign returned subtree to parent's child [OK]

Common Mistakes:

• Not updating parent's child pointer
• Confusing base case with recursive step
• Incorrect comparison operator

5. Suppose the BST can contain duplicate values and you want to find the kth smallest element counting duplicates as separate elements. Which modification to the iterative inorder traversal approach correctly handles duplicates without extra space or preprocessing?

hard

A. Perform a full inorder traversal collecting all values, then remove duplicates before indexing kth element.

B. Modify the traversal to skip nodes with duplicate values to avoid counting duplicates multiple times.

C. Use a hash set to track visited values and only increment count for unique values.

D. Keep the traversal unchanged; duplicates are naturally counted separately in inorder traversal order.

Solution

1. Step 1: Understand duplicates in BST inorder traversal

   Inorder traversal visits nodes in sorted order including duplicates as separate nodes.

2. Step 2: Check if modification is needed

   Since duplicates appear as separate nodes, counting each node visited naturally counts duplicates separately without extra logic.

3. Step 3: Evaluate other options

   Skipping duplicates (B) or using a hash set (C) changes semantics. Removing duplicates after full traversal (D) is inefficient and breaks counting duplicates separately.

4. Final Answer:

   Option D -> Option D

5. Quick Check:

   Inorder traversal inherently counts duplicates separately [OK]

Hint: Duplicates appear as separate nodes in inorder traversal [OK]

Common Mistakes:

• Trying to skip duplicates
• Using extra space unnecessarily
• Assuming duplicates must be merged