Augment a Data Structure

We can break the process of augmenting a data structure into four steps:

  1. Choose an underlying data structure.
  2. Determine additional information to maintain in the underlying data structure.
  3. Verify that we can maintain the additional information for the basic modifying operations on the underlying data structure.
  4. Develop new operations.

Augmenting a data structure involves adding additional information to the basic structure to support more complex queries or updates. By storing this extra data, you can optimize the performance for specific operations. For example, you can augment a binary search tree to maintain the size of each subtree, enabling quick calculations of rank or order statistics.

Java Code to Augment a Binary Search Tree with Subtree Sizes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

class Node {
    int key, size;
    Node left, right;

    public Node(int key) {
        this.key = key;
        this.size = 1;
        this.left = null;
        this.right = null;
    }
}

class AugmentedTree {
    Node root;

    int size(Node node) {
        return (node == null) ? 0 : node.size;
    }

    void updateSize(Node node) {
        if (node != null) {
            node.size = size(node.left) + size(node.right) + 1;
        }
    }

    Node insert(Node node, int key) {
        if (node == null) return new Node(key);

        if (key < node.key) {
            node.left = insert(node.left, key);
        } else {
            node.right = insert(node.right, key);
        }

        updateSize(node);
        return node;
    }
}

C++ Code to Augment a Binary Search Tree with Subtree Sizes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

struct Node {
    int key, size;
    Node *left, *right;
    
    Node(int key): key(key), size(1), left(NULL), right(NULL) {}
};

class AugmentedTree {
    Node* root;
    
public:
    int size(Node* node) {
        return (node == NULL) ? 0 : node->size;
    }

    void updateSize(Node* node) {
        if (node != NULL) {
            node->size = size(node->left) + size(node->right) + 1;
        }
    }

    Node* insert(Node* node, int key) {
        if (node == NULL) return new Node(key);
        
        if (key < node->key) {
            node->left = insert(node->left, key);
        } else {
            node->right = insert(node->right, key);
        }

        updateSize(node);
        return node;
    }
};

Python Code to Augment a Binary Search Tree with Subtree Sizes

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

class Node:
    def __init__(self, key):
        self.key = key
        self.size = 1
        self.left = None
        self.right = None

class AugmentedTree:
    def __init__(self):
        self.root = None

    def size(self, node):
        return 0 if node is None else node.size

    def update_size(self, node):
        if node:
            node.size = self.size(node.left) + self.size(node.right) + 1

    def insert(self, node, key):
        if node is None:
            return Node(key)

        if key < node.key:
            node.left = self.insert(node.left, key)
        else:
            node.right = self.insert(node.right, key)

        self.update_size(node)
        return node

In all implementations, the insert method is responsible for adding a new node to the tree. We then update the size information of each affected node by calling updateSize. With this augmented information, you can perform advanced queries without the need for traversal, thus speeding up certain operations.