Saving Extra information in Self Balanced tree?

Question

In my AVL tree I decided to save the following extra information:

How many nodes in the subtree of the current node

But How can I do this correctly and not to destroy information when doing rotation?

For example my code is:

node *tree::rotate_right(node *ptr) {
    node *tmp_ptr = ptr->left_son;
    ptr->left_son = tmp_ptr->right_son;
    tmp_ptr->right_son = ptr;
    ptr->height = std::max(height(ptr-> left_son), height(ptr-> right_son));
    tmp_ptr->height = std::max(height(tmp_ptr-> left_son), height(tmp_ptr-> right_son));
    return tmp_ptr;
}

node *tree::rotate_left(node *ptr) {
    node *tmp_ptr = ptr->right_son;
    ptr->right_son = tmp_ptr->left_son;
    tmp_ptr->left_son = ptr;
    ptr->height = std::max(height(ptr-> left_son), height(ptr-> right_son));
    tmp_ptr->height = std::max(height(tmp_ptr-> left_son), height(tmp_ptr-> right_son));
    return tmp_ptr;
}

I want this information to find elements by their order, let's say find the i_th element in log(n)

For example: Select(4) here should return 5 using the extra information I saved (I'm not asking for any help with Select implementation but with keeping that extra data)

user_7594711 · Answer

The usual way to do this is to track only the number of left children (plus the total number of nodes as part of the anchor node). Then to find a node at any particular offset you simply go up the tree so long as the current node's count is less than the desired offset. If at any point the parent's count is greater you subtract the current node's count from the offset then you go down the node's right side.

During a left rotate you add the former parent's count to the right child's count, during a right rotate you subtract the left child's count from the former parent's count.

Saving Extra information in Self Balanced tree?

1 Answers