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. You are given a binary search tree and two nodes within it. You need to find the node that is the deepest ancestor common to both nodes, leveraging the BST property to optimize the search. Which approach guarantees the most efficient solution?

easy

A. Perform two separate DFS traversals to find paths from root to each node, then compare paths to find the last common node.

B. Apply dynamic programming to store ancestors of each node and then find the lowest common ancestor by comparing stored data.

C. Iteratively traverse the BST from the root, moving left or right depending on the values of the two nodes relative to the current node, stopping when the current node splits the paths to the two nodes.

D. Use a greedy approach that always moves towards the node with the smaller value until both nodes are found.

Solution

1. Step 1: Understand BST property usage

   In a BST, left subtree nodes are smaller and right subtree nodes are larger than the current node. This allows pruning the search space efficiently.

2. Step 2: Identify optimal traversal

   Starting from root, if both target nodes are smaller, go left; if both are larger, go right; else current node is the split point and thus the LCA.

3. Final Answer:

   Option C -> Option C

4. Quick Check:

   Iterative BST traversal uses O(h) time and O(1) space [OK]

Hint: Use BST property to prune search paths [OK]

Common Mistakes:

• Using path comparison wastes time and space
• Greedy approach ignores split point logic
• Dynamic programming is unnecessary overhead

2. Consider the following code snippet implementing the optimal approach for Two Sum IV in BST. Given the BST with nodes [5,3,6,2,4,null,7] and target k=9, what is the value of variable s during the first iteration of the while loop?

easy

A. 9

B. 5

C. 11

D. 7

Solution

1. Step 1: Perform inorder traversal

   Inorder traversal of BST yields sorted array: [2, 3, 4, 5, 6, 7]

2. Step 2: Calculate s in first iteration

   left=0 (value=2), right=5 (value=7), s = 2 + 7 = 9

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Sum matches target k=9 on first iteration [OK]

Hint: Inorder array first + last sum equals target [OK]

Common Mistakes:

• Off-by-one error in indexing left or right
• Confusing node values with indices
• Miscomputing sum s

3. What is the worst-case time complexity of the optimal iterative BST search algorithm when searching for a value in a BST with n nodes and height h?

medium

A. O(log n) always because BSTs are balanced.

B. O(n) because you might have to visit every node in the tree.

C. O(h) because the search path follows the height of the tree.

D. O(1) because the search uses direct indexing.

Solution

1. Step 1: Understand BST height and search path

   The search path length is at most the height h of the BST, as each step moves down one level.

2. Step 2: Analyze worst-case time complexity

   In worst case (unbalanced tree), h can be up to n, but complexity is expressed as O(h), not O(n) directly.

3. Final Answer:

   Option C -> Option C

4. Quick Check:

   Search complexity depends on tree height, not total nodes necessarily [OK]

Hint: Time depends on tree height, not total nodes always [OK]

Common Mistakes:

• Assuming O(log n) always, ignoring unbalanced trees

4. Suppose you want to extend the BST Iterator to support repeated calls to next() that return the same element multiple times (i.e., allow reuse of elements). Which modification to the Morris traversal based iterator is correct?

hard

A. Remove thread restoration to keep the tree threaded and allow repeated visits to the same node.

B. Store all elements in a list during initialization to allow repeated next() calls without traversal.

C. Modify _advance() to not move current pointer after returning next_val, so next() returns same value again.

D. Use a stack-based iterator and push nodes multiple times to simulate repeated next() calls.

Solution

1. Step 1: Understand reuse requirement

   Allowing repeated next() calls to return the same element requires storing elements since Morris traversal advances and loses previous state.

2. Step 2: Evaluate modifications

   Removing thread restoration (A) corrupts tree, modifying _advance() to not move current (C) breaks traversal logic, and pushing nodes multiple times on stack (B) is inefficient and incorrect. Storing all elements upfront (D) enables repeated next() calls reliably.

3. Final Answer:

   Option B -> Option B

4. Quick Check:

   Pre-storing elements supports repeated next() calls without traversal [OK]

Hint: Reuse requires storing elements, not traversal tricks [OK]

Common Mistakes:

• Assuming traversal can be paused and repeated
• Ignoring tree corruption from missing restoration
• Trying to simulate reuse with threads

5. Suppose you want to modify the BST insertion algorithm to allow inserting duplicate values by always inserting duplicates into the left subtree. Which of the following changes to the iterative insertion code below correctly implements this behavior without breaking the BST property?

def insertIntoBST(root, val):
    if root is None:
        return TreeNode(val)
    curr = root
    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

hard

A. Change the condition to if val <= curr.val: and insert duplicates to the left subtree.

B. Change the condition to if val < curr.val: and insert duplicates to the right subtree.

C. Always insert duplicates as new root nodes to maintain BST property.

D. Disallow duplicates by returning the original tree without insertion.

Solution

1. Step 1: Understand duplicate insertion rule

   Duplicates must go to the left subtree, so values equal to current node should be treated as less or equal.

2. Step 2: Modify comparison operator

   Changing condition to val <= curr.val ensures duplicates follow left subtree insertion path, preserving BST property.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Duplicates inserted left by using <= comparison [OK]

Hint: Use <= to direct duplicates left, preserving BST order [OK]

Common Mistakes:

• Inserting duplicates right breaks problem requirement
• Inserting duplicates as new roots breaks BST
• Ignoring duplicates causes infinite recursion