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. Given the following code snippet for converting a sorted array to a height-balanced BST, what is the value of the root node's value after the first call to helper with input nums = [1, 3, 5]?

easy

A. 3

B. 1

C. 5

D. None

Solution

1. Step 1: Calculate mid index for initial call

   For nums=[1,3,5], left=0, right=2, mid=(0+2)//2=1.

2. Step 2: Identify root value

   nums[mid] = nums[1] = 3, so root node value is 3.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Mid index calculation picks 3 as root [OK]

Hint: Mid index picks middle element as root [OK]

Common Mistakes:

• Choosing first or last element as root
• Off-by-one mid calculation

2. Given the following code for inserting a value into a BST, what will be the value of the variable parent.val when inserting 5 into the BST rooted at 4 with left child 2 and right child 7?

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

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

root = TreeNode(4)
root.left = TreeNode(2)
root.right = TreeNode(7)
val = 5
insertIntoBST(root, val)

easy

A. 4

B. 7

C. 2

D. 5

Solution

1. Step 1: Trace insertion path

   Start at root (4). Since 5 > 4, move right to node 7. Update parent to 4.

2. Step 2: Check right child of 7

   5 < 7, so move left. Left child of 7 is None, so insert here. Parent is now 7, but last parent before insertion was 4.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Parent is last non-null node before insertion -> 4 [OK]

Hint: Parent tracks last node before insertion point [OK]

Common Mistakes:

• Confusing parent with current node
• Off-by-one in loop iteration
• Mistaking inserted value as parent

3. Consider the following buggy code for deleting a node in a BST. Which line contains the subtle bug that can cause the BST property to break when deleting a node with two children?

def deleteNode(root, key):
    if not root:
        return null
    if key < root.val:
        root.left = deleteNode(root.left, key)
    elif key > root.val:
        root.right = deleteNode(root.right, key)
    else:
        if not root.left:
            return root.right
        elif not root.right:
            return root.left
        else:
            successor = findSuccessor(root)
            root.val = successor.val
            root.left = deleteNode(root.right, successor.val)  # Bug here
    return root

medium

A. Line where root.left is assigned deleteNode(root.right, successor.val)

B. Line where root.left is assigned deleteNode(root.left, key)

C. Line where root.val is set to successor.val

D. Line where root.right is assigned deleteNode(root.right, key)

Solution

1. Step 1: Identify two-children deletion logic

   When deleting a node with two children, we replace its value with the successor's value and then delete the successor from the right subtree.

2. Step 2: Spot incorrect subtree assignment

   The line assigns root.left = deleteNode(root.right, successor.val), which incorrectly deletes from the right subtree but assigns result to left child, breaking BST structure.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Correct line should assign root.right = deleteNode(root.right, successor.val) [OK]

Hint: Deleting successor must update right subtree, not left [OK]

Common Mistakes:

• Mixing left and right subtree assignments
• Forgetting to update parent's child pointer
• Replacing value but not deleting successor

4. The following code attempts to validate a BST using inorder traversal. Which line contains a subtle bug that can cause incorrect results on some inputs?

medium

A. Line checking if node is null (if not node: return true)

B. Line updating prev[0] = node.val

C. Line comparing node.val < prev[0] instead of node.val <= prev[0]

D. Line calling inorder(node.left)

Solution

1. Step 1: Identify BST strict inequality

   BST requires node.val to be strictly greater than previous inorder value; using < instead of <= allows duplicates.

2. Step 2: Understand impact of bug

   Using node.val < prev[0] misses the case when node.val == prev[0], incorrectly allowing duplicates and invalid BSTs.

3. Final Answer:

   Option C -> Option C

4. Quick Check:

   Strict inequality check is essential to reject duplicates [OK]

Hint: BST nodes must be strictly increasing in inorder traversal [OK]

Common Mistakes:

• Allowing duplicates by using < instead of <=
• Not resetting prev between calls
• Only comparing immediate children

5. Suppose the BST can contain duplicate values, and you want to find the lowest common ancestor of two nodes p and q that may have the same value. Which modification to the optimal iterative algorithm correctly handles duplicates while maintaining efficiency?

hard

A. Define a consistent rule that duplicates go to the right subtree and use original algorithm without changes.

B. Modify comparisons to use '<=' and '>=' to ensure traversal includes duplicates on both sides.

C. Augment each node with a count of duplicates and traverse left or right based on counts to find LCA.

D. Store paths from root to p and q and compare them to find LCA, since BST property is ambiguous with duplicates.

Solution

1. Step 1: Understand duplicate handling in BST

   BSTs with duplicates require a consistent insertion rule, commonly duplicates go to right subtree.

2. Step 2: Modify comparisons to include duplicates

   Using '<=' and '>=' in comparisons ensures traversal includes duplicates correctly on both sides.

3. Step 3: Avoid ambiguous or inefficient methods

   Path comparison or augmenting nodes adds overhead and complexity.

4. Final Answer:

   Option B -> Option B

5. Quick Check:

   Inclusive comparisons handle duplicates while preserving efficiency [OK]

Hint: Modify comparison operators to handle duplicates correctly [OK]

Common Mistakes:

• Using strict comparisons that miss duplicates
• Reverting to path comparison losing efficiency
• Adding unnecessary data structures