← Back to DevBytes

Solving Kth Smallest Element in BST in JavaScript: Step-by-Step Guide

Solving Kth Smallest Element in BST in JavaScript: Step-by-Step Guide

Binary Search Trees (BSTs) are one of the most foundational data structures in computer science. Among the many problems you can solve with them, finding the Kth smallest element is a classic interview question that tests your understanding of tree traversal, recursion, and optimization. In this tutorial, we'll walk through everything you need to know to solve this problem efficiently in JavaScript.

What Is the Kth Smallest Element Problem?

Given a Binary Search Tree and an integer k, the task is to return the kth smallest element (1-indexed) among all the nodes in the tree. For example, if k = 1, you return the smallest element; if k = 3, you return the third smallest.

The key insight that makes this problem tractable is the defining property of a BST: for any node, all values in its left subtree are smaller, and all values in its right subtree are larger. This means an inorder traversal of a BST visits nodes in ascending order.

Defining the Tree Node

Before diving into solutions, let's define the basic building block — the tree node:

class TreeNode {
  constructor(val, left = null, right = null) {
    this.val = val;
    this.left = left;
    this.right = right;
  }
}

Here's a sample BST we'll use throughout this tutorial:

// Constructing the BST:
//        5
//       / \
//      3   7
//     / \ / \
//    1  4 6  8

const root = new TreeNode(5,
  new TreeNode(3,
    new TreeNode(1),
    new TreeNode(4)
  ),
  new TreeNode(7,
    new TreeNode(6),
    new TreeNode(8)
  )
);

The inorder traversal of this tree produces [1, 3, 4, 5, 6, 7, 8]. So the 3rd smallest element is 4, and the 5th smallest is 6.

Why This Problem Matters

This problem is more than an academic exercise. It has real-world applications and teaches several important concepts:

Approach 1: Recursive Inorder Traversal

The most intuitive approach leverages the fact that an inorder traversal of a BST yields nodes in sorted order. We can perform an inorder traversal, collect the values into an array, and then return the element at index k - 1.

Implementation

function kthSmallestRecursive(root, k) {
  const result = [];

  function inorder(node) {
    if (node === null) return;

    inorder(node.left);
    result.push(node.val);
    inorder(node.right);
  }

  inorder(root);
  return result[k - 1];
}

// Usage:
console.log(kthSmallestRecursive(root, 3)); // Output: 4
console.log(kthSmallestRecursive(root, 5)); // Output: 6

Analysis

This approach is clean and easy to understand, but it has a drawback: it always traverses the entire tree, even if k is small. The time complexity is O(n) and the space complexity is O(n) for the result array plus O(h) for the recursion stack, where h is the height of the tree.

Approach 2: Optimized Recursive Inorder with Early Termination

We can improve the recursive approach by stopping the traversal as soon as we've found the kth element. Instead of collecting all values, we maintain a counter and return early.

Implementation

function kthSmallestOptimized(root, k) {
  let count = 0;
  let answer = null;

  function inorder(node) {
    if (node === null || answer !== null) return;

    inorder(node.left);

    count++;
    if (count === k) {
      answer = node.val;
      return;
    }

    inorder(node.right);
  }

  inorder(root);
  return answer;
}

// Usage:
console.log(kthSmallestOptimized(root, 1)); // Output: 1
console.log(kthSmallestOptimized(root, 7)); // Output: 8

Analysis

This version stops traversing once the answer is found. In the best case (when k is small), it runs in O(k) time. In the worst case, it's still O(n). The space complexity is O(h) for the recursion stack, which is O(log n) for a balanced tree and O(n) for a completely skewed tree.

Approach 3: Iterative Inorder Traversal

Recursive solutions can cause stack overflow errors on very deep trees. An iterative approach using an explicit stack avoids this issue and gives us even finer control over early termination.

Implementation

function kthSmallestIterative(root, k) {
  const stack = [];
  let current = root;

  while (current !== null || stack.length > 0) {
    // Go as far left as possible
    while (current !== null) {
      stack.push(current);
      current = current.left;
    }

    // Process the next node in order
    current = stack.pop();
    k--;

    if (k === 0) {
      return current.val;
    }

    // Move to the right subtree
    current = current.right;
  }

  return null; // k is larger than the number of nodes
}

// Usage:
console.log(kthSmallestIterative(root, 2)); // Output: 3
console.log(kthSmallestIterative(root, 6)); // Output: 7

How It Works

The iterative approach simulates the call stack manually. We push all left children onto the stack until we reach a leaf, then pop and process nodes one at a time. After processing a node, we move to its right subtree and repeat. By decrementing k each time we process a node, we know we've found our answer when k reaches zero.

This approach has the same time and space complexity as the optimized recursive version but avoids recursion depth issues.

Approach 4: Augmented BST for Repeated Queries

If you need to perform many kth-smallest queries on the same tree, traversing each time is wasteful. Instead, you can augment each node to store the size of its left subtree. This lets you navigate directly to the kth element in O(h) time per query.

Implementation

class AugmentedTreeNode {
  constructor(val, left = null, right = null) {
    this.val = val;
    this.left = left;
    this.right = right;
    this.leftCount = 0; // Number of nodes in the left subtree
  }
}

function buildAugmentedTree(val, left = null, right = null) {
  const node = new AugmentedTreeNode(val, left, right);
  node.leftCount = countNodes(left);
  return node;
}

function countNodes(node) {
  if (node === null) return 0;
  return 1 + countNodes(node.left) + countNodes(node.right);
}

function kthSmallestAugmented(root, k) {
  let current = root;

  while (current !== null) {
    const leftSize = current.leftCount;

    if (k === leftSize + 1) {
      return current.val;
    } else if (k <= leftSize) {
      current = current.left;
    } else {
      k = k - leftSize - 1;
      current = current.right;
    }
  }

  return null;
}

// Usage:
const aRoot = buildAugmentedTree(5,
  buildAugmentedTree(3,
    buildAugmentedTree(1),
    buildAugmentedTree(4)
  ),
  buildAugmentedTree(7,
    buildAugmentedTree(6),
    buildAugmentedTree(8)
  )
);

console.log(kthSmallestAugmented(aRoot, 3)); // Output: 4
console.log(kthSmallestAugmented(aRoot, 4)); // Output: 5

Analysis

Building the augmented tree takes O(n) time initially, but each subsequent query runs in O(h) time — O(log n) for a balanced tree. The trade-off is additional space (O(n) for the leftCount fields) and the need to maintain these counts during insertions and deletions.

Handling Edge Cases

A robust solution must handle several edge cases gracefully:

Here's a wrapper function that validates inputs:

function kthSmallestSafe(root, k) {
  if (root === null) {
    throw new Error("Tree is empty");
  }
  if (k < 1) {
    throw new Error("k must be at least 1");
  }

  const result = kthSmallestIterative(root, k);

  if (result === null) {
    throw new Error("k is larger than the number of nodes in the tree");
  }

  return result;
}

Best Practices

Choose the Right Approach for Your Use Case

If you're solving this once on a given tree, the iterative approach offers the best balance of efficiency and clarity. If you need repeated queries, invest in the augmented BST approach. For quick prototyping or when code simplicity matters most, the basic recursive solution is perfectly acceptable.

Keep the Tree Balanced

A skewed BST degenerates into a linked list, making all operations O(n) instead of O(log n). Consider using self-balancing trees like AVL trees or Red-Black trees if you're building the tree from scratch and performance is critical.

Avoid Global Mutable State

In the optimized recursive approach, we used variables declared outside the inner function. While convenient, this pattern can cause bugs if the function is called multiple times or in concurrent contexts. An alternative is to return values up the call chain:

function kthSmallestPure(root, k) {
  function helper(node, k) {
    if (node === null) return { found: false, value: null, count: 0 };

    // Check left subtree
    const left = helper(node.left, k);
    if (left.found) return left;

    // Check current node
    const currentCount = left.count + 1;
    if (currentCount === k) {
      return { found: true, value: node.val, count: currentCount };
    }

    // Check right subtree
    const right = helper(node.right, k - currentCount);
    if (right.found) return right;

    return { found: false, value: null, count: currentCount + right.count };
  }

  const result = helper(root, k);
  return result.value;
}

Write Tests

Always test your implementation against various tree shapes and values:

function runTests() {
  // Test 1: Balanced tree
  const tree1 = new TreeNode(5,
    new TreeNode(3, new TreeNode(1), new TreeNode(4)),
    new TreeNode(7, new TreeNode(6), new TreeNode(8))
  );
  console.assert(kthSmallestIterative(tree1, 1) === 1, "Test 1.1 failed");
  console.assert(kthSmallestIterative(tree1, 4) === 5, "Test 1.2 failed");
  console.assert(kthSmallestIterative(tree1, 7) === 8, "Test 1.3 failed");

  // Test 2: Single node
  const tree2 = new TreeNode(42);
  console.assert(kthSmallestIterative(tree2, 1) === 42, "Test 2.1 failed");

  // Test 3: Left-skewed tree
  const tree3 = new TreeNode(3, new TreeNode(2, new TreeNode(1)));
  console.assert(kthSmallestIterative(tree3, 1) === 1, "Test 3.1 failed");
  console.assert(kthSmallestIterative(tree3, 3) === 3, "Test 3.2 failed");

  // Test 4: Right-skewed tree
  const tree4 = new TreeNode(1, null, new TreeNode(2, null, new TreeNode(3)));
  console.assert(kthSmallestIterative(tree4, 2) === 2, "Test 4.1 failed");

  console.log("All tests passed!");
}

runTests();

Comparing the Approaches

Here's a quick summary to help you choose:

Conclusion

Finding the kth smallest element in a BST is a problem that beautifully illustrates the power of inorder traversal and the importance of choosing the right algorithm for your specific constraints. The iterative approach is generally the best default choice, offering early termination without the risk of stack overflow. For scenarios involving repeated queries, the augmented BST approach pays for its preprocessing cost with lightning-fast lookups. By understanding all four approaches — from the straightforward recursive solution to the optimized augmented tree — you'll be well-equipped to handle this problem in interviews and real-world applications alike. Remember to always consider edge cases, keep your trees balanced when possible, and write thorough tests to validate your implementation.

— Ad —

Google AdSense will appear here after approval

← Back to all articles