Finding the Kth Largest Element in a Binary Search Tree

Problem Statement

Given a binary search tree (BST) and a positive integer K, the task is to find the Kth largest element in the BST. The Kth largest element is the Kth largest value when all elements in the tree are sorted in descending order.

Introduction

The problem of finding the Kth largest element in a Binary Search Tree is a common task in computer science and has applications in various domains such as databases, statistics, finance, and more. In this blog post, we'll explore the problem and provide a generalized solution in three popular programming languages: Java, Python, and C++.

Overview

A Binary Search Tree is a data structure that allows for efficient searching, insertion, and deletion of elements. To find the Kth largest element in a BST, we can perform a reverse in-order traversal of the tree while keeping track of the count of visited nodes. When the count reaches K, we have found the Kth largest element.

Code

Java Implementation

/* Copyrights to venkys.io
For more information, visit https://venkys.io */

//Space complexity: O(H)
//Time complexity: O(H + K)

import java.util.Scanner;

class TreeNode {
    int val;
    TreeNode left, right;

    public TreeNode(int val) {
        this.val = val;
        this.left = this.right = null;
    }
}

class KthLargestInBST {
    private int count = 0;
    private int result = -1;

    /**
     * Finds the kth largest element in a Binary Search Tree (BST).
     *
     * root -- The root of the BST.
     *  k -- The value of k for finding the kth largest element.
     * return The kth largest element in the BST.
     */
    public int kthLargest(TreeNode root, int k) {
        count = k;
        inorderTraversal(root);
        return result;
    }

    /**
     * Performs an inorder traversal of the BST to find the kth largest element.
     *
     * node -- The current node during the traversal.
     */
    private void inorderTraversal(TreeNode node) {
        if (node == null || count == 0) {
            return;
        }

        inorderTraversal(node.right);
        count--;

        if (count == 0) {
            result = node.val;
            return;
        }

        inorderTraversal(node.left);
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);

        // Input the values for the tree nodes in a space-separated format (e.g., "5 3 7 2 -1 6 8").
        String[] nodeValues = scanner.nextLine().split(" ");
        TreeNode root = buildTree(nodeValues, 0);

        KthLargestInBST kthLargestInBST = new KthLargestInBST();

        // Input the value of k.
        // System.out.print("Enter the value of k: ");
        int k = scanner.nextInt();

        int result = kthLargestInBST.kthLargest(root, k);
        System.out.println(result);
    }

    private static TreeNode buildTree(String[] nodeValues, int index) {
        if (index >= nodeValues.length || nodeValues[index].equals("-1")) {
            return null;
        }

        TreeNode root = new TreeNode(Integer.parseInt(nodeValues[index]));
        root.left = buildTree(nodeValues, 2 * index + 1);
        root.right = buildTree(nodeValues, 2 * index + 2);

        return root;
    }
}

Python Implementation

""" Copyrights to venkys.io
For more information, visit https://venkys.io

Space complexity: O(H)
Time complexity: O(H + K) """

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

class KthLargestInBST:
    def __init__(self):
        self.count = 0
        self.result = None

    def kth_largest(self, root, k):
        self.count = k
        self.inorder_traversal(root)
        return self.result

    def inorder_traversal(self, node):
        if not node or self.count == 0:
            return

        self.inorder_traversal(node.right)

        self.count -= 1
        if self.count == 0:
            self.result = node.val
            return

        self.inorder_traversal(node.left)

def build_tree_from_input(nodes, index):
    if index < len(nodes) and nodes[index] is not None:
        root = TreeNode(nodes[index])
        root.left = build_tree_from_input(nodes, 2 * index + 1)
        root.right = build_tree_from_input(nodes, 2 * index + 2)
        return root
    return None

if __name__ == "__main__":
    # Input the values for the tree nodes (e.g., "5 3 7 2 null 6 8")
    nodes_input = input()
    nodes = [int(val) if val != "null" else None for val in nodes_input.split()]

    root = build_tree_from_input(nodes, 0)

    if root is None:
        print("The BST is empty.")
    else:
        kth_largest_obj = KthLargestInBST()

        # Input the value of k
        k = int(input("Enter the value of k: "))

        result = kth_largest_obj.kth_largest(root, k)
        print(result)

C++ Implementation

/* Copyrights to venkys.io
For more information, visit https://venkys.io

Space complexity: O(H)
Time complexity: O(H + K) */

#include <iostream>
#include <cstdlib>  // For using NULL
#include <vector>
#include <sstream>
using namespace std;

// Definition of a TreeNode in a Binary Search Tree (BST)
struct TreeNode {
    int val;
    TreeNode* left;
    TreeNode* right;

    // Constructor to initialize a TreeNode with a given value
    TreeNode(int val) : val(val), left(nullptr), right(nullptr) {}
};

// Class to find the kth largest element in a BST
class KthLargestInBST {
private:
    int count;   // Counter to keep track of the number of nodes visited
    int result;  // Variable to store the result (kth largest element)

public:
    // Constructor to initialize the count and result
    KthLargestInBST() : count(0), result(-1) {}

    // Function to find the kth largest element in the BST
    int kthLargest(TreeNode* root, int k);

private:
    // Helper function for inorder traversal to find kth largest element
    void inorderTraversal(TreeNode* node);
};

// Function to build a BST from values in a vector
TreeNode* buildBST(const vector<int>& values, size_t& index);

int main() {
    // Prompt the user to enter values for BST nodes and the value of k
    // cout << "Enter the values for the BST nodes and the value of k (in a space-separated format):" << endl;

    // Read input line
    string inputLine;
    getline(cin, inputLine);

    // Parse values from the input line
    stringstream ss(inputLine);
    vector<int> values;
    int val;
    while (ss >> val) {
        values.push_back(val);
    }

    // Check if there are enough values to build a BST and get the value of k
    if (values.size() < 2) {
        cout << "Invalid input format. Please enter values for BST nodes and the value of k." << endl;
        return 1;
    }

    // Get the value of k
    int k = values.back();
    values.pop_back();

    // Build the Binary Search Tree (BST) based on the provided values
    size_t index = 0;
    TreeNode* root = buildBST(values, index);

    if (root == NULL) {
        cout << "The BST is empty." << endl;
        return 0;
    }

    // Create an instance of the KthLargestInBST class
    KthLargestInBST kthLargestFinder;

    // Find and display the kth largest element in the BST
    int result = kthLargestFinder.kthLargest(root, k);

    if (result != -1) {
        cout << result << endl;
    } else {
        cout << "Invalid value of k or the BST is not large enough." << endl;
    }

    return 0;
}

// Implementation of the kthLargest function from KthLargestInBST class
int KthLargestInBST::kthLargest(TreeNode* root, int k) {
    if (root == NULL || k <= 0) {
        return -1;  // Invalid input or empty tree
    }

    // Set count to k and perform inorder traversal to find the kth largest element
    count = k;
    inorderTraversal(root);
    return result;
}

// Implementation of the inorderTraversal function for BST traversal
void KthLargestInBST::inorderTraversal(TreeNode* node) {
    // Base case: If the node is null or kth largest element is already found, return
    if (node == NULL || count == 0) {
        return;
    }

    // Traverse the right subtree (reverse order)
    inorderTraversal(node->right);

    // Process the current node
    if (--count == 0) {
        result = node->val;  // Update result when kth largest element is found
        return;
    }

    // Traverse the left subtree
    inorderTraversal(node->left);
}

// Implementation of the buildBST function to construct a BST from values in a vector
TreeNode* buildBST(const vector<int>& values, size_t& index) {
    // Base case: Return null if the index is out of bounds
    if (index >= values.size()) {
        return nullptr;
    }

    // Read the current value from the vector
    int val = values[index++];

    // Return null if the value is -1 (null node)
    if (val == -1) {
        return nullptr;
    }

    // Create a new TreeNode with the current value
    TreeNode* node = new TreeNode(val);

    // Recursively build left and right subtrees
    node->left = buildBST(values, index);
    node->right = buildBST(values, index);

    return node;
}

Step-by-Step Explanation

Node Definition: Define a TreeNode class to represent a node in the binary tree.
Class Initialization: Create a class (KthLargestInBST) to encapsulate the logic for finding the Kth largest element.
Method for Kth Largest: Implement a method (kthLargest in Java and Python, and kthLargest in C++) that takes the root of the tree and the value of K as parameters. It initializes a count and calls an inorderTraversal method.
In-Order Traversal: Implement the inorderTraversal method that performs a reverse in-order traversal. It starts from the right, decrements the count, and stores the value when the count reaches 0. It then moves to the left.

Time and Space Complexity Analysis

Time Complexity: The time complexity of the solution is O(H + K), where H is the height of the tree and K is the Kth element to be found. In the worst case, when K is small compared to the total number of nodes, it doesn't require traversing the entire tree.
Space Complexity: The space complexity is O(H), where H is the height of the tree, due to the recursive calls on the stack.

Real-World Applications

Efficiently finding the Kth largest element in a BST has real-world applications in various domains:

Databases: In databases, indexing and searching for specific data are common tasks. Efficiently finding the Kth largest value can be useful for optimizing database queries and report generation.
Statistics: Data analysis often involves finding the top K values in a dataset. This algorithm can be applied to efficiently compute percentiles and other statistical measures.
Financial Analysis: In financial markets, finding the Kth largest stock price or asset value is essential for various investment strategies and risk assessment.
Medicine: In medical research, analyzing patient data might require identifying the Kth largest value, such as the highest recorded temperature or blood pressure, to monitor health trends.

The ability to efficiently find the Kth largest element in a BST is a valuable tool in these and many other fields where data analysis and decision-making rely on ordered data.

Test Cases

Test Case 1: Input: 5 3 7 2 null 6 8 3 Output: 6

Explanation: The provided input constructs a Binary Search Tree (BST) with the following structure:

We're asked to find the 3rd largest element in the BST.
In-order traversal of the BST would give us elements in ascending order.
So, the 3rd largest element would be the third element from the end of the in-order traversal sequence.
In this case, the in-order traversal sequence is [2, 3, 5, 6, 7, 8], and the 3rd largest element is 6.
Therefore, the output is 6, indicating the 3rd largest element in the BST.

Test Case 2: Input: 8 5 10 3 7 9 12 4 Output: 9

Explanation: The provided input constructs a Binary Search Tree (BST) with the following structure:

     8
    / \
   5   10
  / \  / \
 3   7 9  12

We're asked to find the 4th largest element in the BST.
In-order traversal of the BST would give us elements in ascending order.
So, the 4th largest element would be the fourth element from the end of the in-order traversal sequence.
In this case, the in-order traversal sequence is [3, 5, 7, 8, 9, 10, 12], and the 4th largest element is 9.
Therefore, the output is 9, indicating the 4th largest element in the BST.