A Problem using Augmentation of Data Structure

Given a set of integers, the task is to design a data structure that supports two operations efficiently:

Insertion: Insert an integer into the data structure.
Query: Given an integer k, find the number of elements in the data structure that are less than or equal to k.

Naive Approach: A simple way to solve this problem is to use an array or a list to store the integers. For insertion, you can just append the integer to the end of the list. For the query operation, iterate through the list and count the elements less than or equal to k. This approach takes linear time for the query operation.

Augmenting the Data Structure: To improve the efficiency of the query operation, we can augment the data structure with additional information. One way to do this is by maintaining a sorted order of elements. This can be achieved using a balanced Binary Search Tree (BST).

Structure of Augmented Tree: Each node in the BST should store the value of the element, the size of the subtree rooted at that node, and pointers to the left and right children.

Insertion:
- Insert the element into the BST as you normally would.
- While inserting, update the size of each visited node.
Query:
- Start at the root of the BST.
- Traverse the tree:
  - If the current node’s value is less than or equal to k, add the size of its left subtree (including itself) to the count.
  - Move to the left or right child accordingly.

Implementation of above algorithm is given below:

C++

// C++ Implementation
 
#include <iostream>
 
// Class to define node for the Augmented BST
class AugmentedBSTNode {
public:
    int value;
    int size;
    AugmentedBSTNode* left;
    AugmentedBSTNode* right;
 
    AugmentedBSTNode(int value) {
        this->value = value;
        this->size = 1;
        this->left = nullptr;
        this->right = nullptr;
    }
};
 
// Class to define the augmented BST with the operations
class AugmentedBST {
private:
    AugmentedBSTNode* root;
 
    // Method to insert a value in the BST
    AugmentedBSTNode* _insert(AugmentedBSTNode* node, int value) {
        if (node == nullptr) {
            return new AugmentedBSTNode(value);
        }
        node->size++;
        if (value <= node->value) {
            node->left = _insert(node->left, value);
        } else {
            node->right = _insert(node->right, value);
        }
        return node;
    }
 
    // Helper Method to query in the BST
    int _query(AugmentedBSTNode* node, int k) {
        if (node == nullptr) {
            return 0;
        }
        if (node->value <= k) {
            return (node->left != nullptr ? node->left->size : 0) + 1 + _query(node->right, k);
        } else {
            return _query(node->left, k);
        }
    }
 
public:
    AugmentedBST() {
        this->root = nullptr;
    }
 
    void insert(int value) {
        root = _insert(root, value);
    }
 
    int query(int k) {
        return _query(root, k);
    }
};
 
// Driver Code
int main() {
    AugmentedBST augmentedDS;
    augmentedDS.insert(5);
    augmentedDS.insert(2);
    augmentedDS.insert(8);
    augmentedDS.insert(1);
 
    int result = augmentedDS.query(6);
    std::cout << result << std::endl;
 
    return 0;
}
 
// This code is contributed by Tapesh(tapeshdu420)

Java

// Class to define node for the Augmented BST
class AugmentedBSTNode {
    int value;
    int size;
    AugmentedBSTNode left;
    AugmentedBSTNode right;
 
    public AugmentedBSTNode(int value) {
        this.value = value;
        this.size = 1;
        this.left = null;
        this.right = null;
    }
}
 
// Class to define the augmented BST with the operations
class AugmentedBST {
    private AugmentedBSTNode root;
 
    // Method to insert a value in the BST
    public void insert(int value) {
        root = _insert(root, value);
    }
 
    // Helper Method to insert a node in the Augmented BST
    private AugmentedBSTNode _insert(AugmentedBSTNode node, int value) {
        if (node == null) {
            return new AugmentedBSTNode(value);
        }
        node.size++;
        if (value <= node.value) {
            node.left = _insert(node.left, value);
        } else {
            node.right = _insert(node.right, value);
        }
        return node;
    }
 
    // Method to query in the BST
    public int query(int k) {
        return _query(root, k);
    }
 
    // Helper Method to query in the BST
    private int _query(AugmentedBSTNode node, int k) {
        if (node == null) {
            return 0;
        }
        if (node.value <= k) {
            return (node.left != null ? node.left.size : 0) + 1 + _query(node.right, k);
        } else {
            return _query(node.left, k);
        }
    }
}
 
// Driver Code
public class Main {
    public static void main(String[] args) {
        AugmentedBST augmentedDS = new AugmentedBST();
        augmentedDS.insert(5);
        augmentedDS.insert(2);
        augmentedDS.insert(8);
        augmentedDS.insert(1);
 
        int result = augmentedDS.query(6);
        System.out.println(result);
    }
}

Python

# class to define node for the Augmented BST
class AugmentedBSTNode:
    def __init__(self, value):
        self.value = value
        self.size = 1
        self.left = None
        self.right = None
 
# class to define the augmented BST with the operations
class AugmentedBST:
    def __init__(self):
        self.root = None
 
    # Method to insert a value in the BST
    def insert(self, value):
        self.root = self._insert(self.root, value)
 
    # Helper Method to insert a node in the Augmented BST
    def _insert(self, node, value):
        if not node:
            return AugmentedBSTNode(value)
        node.size += 1
        if value <= node.value:
            node.left = self._insert(node.left, value)
        else:
            node.right = self._insert(node.right, value)
        return node
 
    # Method to query in the BST
    def query(self, k):
        return self._query(self.root, k)
 
    # Helper Method to query in the BST
    def _query(self, node, k):
        if not node:
            return 0
        if node.value <= k:
            return (node.left.size if node.left else 0) + 1 + self._query(node.right, k)
        else:
            return self._query(node.left, k)
 
 
# Example usage:
augmented_ds = AugmentedBST()
augmented_ds.insert(5)
augmented_ds.insert(2)
augmented_ds.insert(8)
augmented_ds.insert(1)
 
result = augmented_ds.query(6)
print(result)

C#

using System;
 
// Class to define node for the Augmented BST
class AugmentedBSTNode
{
    public int value;
    public int size;
    public AugmentedBSTNode left;
    public AugmentedBSTNode right;
 
    public AugmentedBSTNode(int value)
    {
        this.value = value;
        this.size = 1;
        this.left = null;
        this.right = null;
    }
}
 
// Class to define the augmented BST with the operations
class AugmentedBST
{
    private AugmentedBSTNode root;
 
    // Method to insert a value in the BST
    private AugmentedBSTNode Insert(AugmentedBSTNode node, int value)
    {
        if (node == null)
        {
            return new AugmentedBSTNode(value);
        }
        node.size++;
        if (value <= node.value)
        {
            node.left = Insert(node.left, value);
        }
        else
        {
            node.right = Insert(node.right, value);
        }
        return node;
    }
 
    // Helper Method to query in the BST
    private int Query(AugmentedBSTNode node, int k)
    {
        if (node == null)
        {
            return 0;
        }
        if (node.value <= k)
        {
            return (node.left != null ? node.left.size : 0) + 1 + Query(node.right, k);
        }
        else
        {
            return Query(node.left, k);
        }
    }
 
    public AugmentedBST()
    {
        this.root = null;
    }
 
    public void Insert(int value)
    {
        root = Insert(root, value);
    }
 
    public int Query(int k)
    {
        return Query(root, k);
    }
}
 
// Driver Code
class Program
{
    static void Main()
    {
        AugmentedBST augmentedDS = new AugmentedBST();
        augmentedDS.Insert(5);
        augmentedDS.Insert(2);
        augmentedDS.Insert(8);
        augmentedDS.Insert(1);
 
        int result = augmentedDS.Query(6);
        Console.WriteLine(result);
    }
}

Javascript

// Class to define node for the Augmented BST
class AugmentedBSTNode {
    constructor(value) {
        this.value = value;
        this.size = 1;
        this.left = null;
        this.right = null;
    }
}
 
// Class to define the augmented BST with operations
class AugmentedBST {
    constructor() {
        this.root = null;
    }
 
    // Method to insert a value in the BST
    insert(value) {
        this.root = this._insert(this.root, value);
    }
 
    // Helper Method to insert a node in the Augmented BST
    _insert(node, value) {
        if (node === null) {
            return new AugmentedBSTNode(value);
        }
        node.size++;
        if (value <= node.value) {
            node.left = this._insert(node.left, value);
        } else {
            node.right = this._insert(node.right, value);
        }
        return node;
    }
 
    // Method to query in the BST
    query(k) {
        return this._query(this.root, k);
    }
 
    // Helper Method to query in the BST
    _query(node, k) {
        if (node === null) {
            return 0;
        }
        if (node.value <= k) {
            return (node.left !== null ? node.left.size : 0) + 1 + this._query(node.right, k);
        } else {
            return this._query(node.left, k);
        }
    }
}
 
// Driver Code
const augmentedDS = new AugmentedBST();
augmentedDS.insert(5);
augmentedDS.insert(2);
augmentedDS.insert(8);
augmentedDS.insert(1);
 
const result = augmentedDS.query(6);
console.log(result);

Output

Time Complexity: O(Q * log(N)), where Q is the number of queries and N is the number of nodes in the BST. After Augmenting the BST, each query can be answered in O(logN) time.
Auxiliary Space: O(N)

A Problem using Augmentation of Data Structure

C++

Java

Python

C#

Javascript

Introduction to Augmented Data Structure

Categories

Contact US

A Problem using Augmentation of Data Structure

C++

Java

Python

C#

Javascript

Introduction to Augmented Data Structure

Similar Reads

Categories

Contact US