Augmentation (Data Structures)

Augmentation is a problem solving technique where we take an existing data structure which gets us ~80% of the way to the solution and "augment" it with additional information which supports the new/changed operations we need.

Example: Ordered Set

An ordered set is a collection of unique elements where each element has a "rank" which represents its ordering relative to the other elements of the collection. If an ordered set has n elements, the minimal element has rank 1 and the maximal element has rank n.

An ordered set should support the standard insert/remove/search operations, but should also support retrieving the rank of an element as well as retrieving the element with the given rank.

We can implement this by augmenting an AVL tree such that each node also stores the size of its sub-tree. Therefore if we maintain a rank parameter which is initialized to 0 as we traverse the tree, the rank of any node is given by the size of its left sub-tree + 1 provided we update rank to equal the size of its left sub-tree + 1 when we traverse into the right sub-tree.

rank(root, x):
  rank = 0
  while root != nil:
    if x < root:
      root = root.left
    else if x > root:
      rank += root.left.size + 1
      root = root.right
    else:
      return rank + root.left.size + 1
  return nil

This supports the implementation of retrieving the rank of an element, but in order to retrieve an element by rank, we use the following procedure

select(root, rank):
  if rank < root.left.size + 1:
    return select(root.left, rank)
  else if rank > root.left.size + 1:
    return select(root.right, rank - (root.left.size + 1))
  else:
    return root

In this case, we're only concerned with finding the element, not keeping track of the actual rank of the element, so we can get away with subtracting the rank of each node during each recursive step which traverses the right child.

In order to keep the size attribute correct, we update it for each node in the tree we visit during the insert/delete operations.

Procedure

Use the following steps to solve a problem using augmentation:

1. Pick the base data structure
2. Define the additional information
3. Make sure the desirable properties of the original data structure are not hindered by the addition
4. Define the additional operations