How to find the index of an element in sorted set?

I can find an element in a sorted set (backed by BST) in O(logN). Now I would like the index of this element. For example, in set {1, 3, 4, 10}, the index of 4 is 2 and the index of 1 is 0.

Obviously, I can just iterate over the set, so the trivial solution is O(N). Can we do better using probably BST properties and/or auxiliary data structures ? 开发者_如何转开发

With a simple BST where the order of insertion of elements in random, you have way to determine how many elements exactly are smaller than a given element without walking the tree.

If you had a balanced tree, such as a red-black tree, then you could at least put a lower and upper bound on the index due to the bounds on the height of the tree. If the order of insertion of elements to a BST is non-random then again, you could say something about the tree height without walking it and give some estimate of the approximate index.

As for auxiliary data structures, you could create an auxiliary dictionary that maps elements to their index. However, building that index takes O(N) and the index becomes stale when you add new elements to the BST, so this only works well for BSTs with infrequent updates.

Another solution is to extend the BST nodes with two properties: index and count. The index says how many elements smaller than the one in this node are in the tree. The count says how many elements were in the BST when you last updated that node's index. With relatively simple changes to insertion, deletion and search on the BST, that don't affect those basic operations beyond a constant time, and can get the element's index directly in O(1).

Essentially, as you insert a new node, for every node you pass on your path downward, if the new element is smaller (i.e. your next step is to the left child), increment both index and count of that node. When you find the new element's place, you give it a count based on its parent, and an index based on its parent and left child. This leaves the elements larger than the new one with a wrong index, but you could easily update that as you search for an element by referring to the count value of the parent - the difference between the count of the parent and the child tells you how many insertions of smaller elements happened since you last updated the child's index, so you simply add that difference to the index.