如何利用红黑树实现排名?

1. 问题

红黑树是一种自平衡的二叉查找树，它可以在O(logn)时间内执行查找、插入和删除。在c++ STL，linux内核中都有使用。

红黑树本身是有序的，现在问题是对于指定的元素，如何能快速查到它在整个元素集的排名，或者根据排名快速查询对应的元素？

2. 思路

排名分顺序和逆序，这里只讨论顺序的情况。顺序的话排名就是求比当前元素小的元素的个数，根据红黑树的性质，左子树的节点都比根节点小，右子树的节点都比根节点大，求排名就等价于求节点左子树元素的个数。

根据树的递归性质，我们只需要在每个节点增加一个字段count用来统计当前节点子树的个数，同时在红黑树做插入、删除操作的时候更新count字段，就能在O(logn)的时间内查询到该元素的排名。

3. 实现

红黑树节点增加count字段，count[x]表示x节点子节点元素的个数，包括它的左子树，它的右子树和它自己本身。

count[x] = count[left[x]] + count[right[x]] + 1; // x非空
count[x] = 0; // x为空

红黑树旋转的时候，保证count满足我们的定义就可以。

3.1 左旋

如图所示，左旋的时候，只有x和y节点count域会发生变化，只需要修改图中x和y节点对应的count，其余节点不需要修改。

左旋后：

count[x] = count[α] + count[β] + 1
count[y] = count[x] + count[γ] + 1

左旋伪代码:

LEFT-ROTATE(T, x)
    y = right[x]
    right[x] = left[y]
    p[left[y]] = x
    p[y] = p[x]
    count[x] = count[left[x]] + count[left[right[x]]] + 1
    count[right[x]] = count[left[x]] + count[left[right[x]]] + count[right[right[x]]] + 2
    if p[x] == nil
    then root[T] = y
    else if x == left[p[x]]
        then left[p[x]] = y
        else right[p[x]] = y
    left[y] = x
    p[x] = y

3.2 右旋

右旋同样类似于左旋，只需要修改x 和 y节点对应的count。

右旋后：

count[y] = count[γ] + count[β] + 1
count[x] = count[x] + count[α] + 1

右旋伪代码：

RIGHT-ROTATE(T, y)
    x = left[y]
    left[y] = right[x]
    p[right[x]] = y
    p[x] = p[y]
    count[y] = count[right[y]] + count[right[left[y]]] + 1
    count[left[y]] = count[right[y]] + count[left[left[y]]] + count[right[left[y]]] + 2
    if p[y] == nil
    then root[T] = x
    else if y == right[p[y]]
        then right[p[y]] = x
        else left[p[y]] = x
    right[x] = y
    p[y] = x

插入和删除的时候对于count的修改比较简单，只修改节点所有祖先节点的count，插入的时候，我们先按照红黑树的规则插入到指定位置，然后对该节点的所有祖先节点的count都增加1，然后再做平衡调整，删除的时候类似。

3.3 根据排名查询元素

跟红黑树普通的查询类似，只不过用来比较的域换成了count，这里分为三种情况：

1.节点左子树个数 + 1 == rank，表示已经找到需要查询的元素

2.节点左子树个数 + 1 > rank, 表示当前节点左子树个数大于rank - 1，我们需要在左子树中递归查询

3.节点左子树个数 + 1 < rank, 表示当前节点左子树个数大于rank - 1, 我们需要在右字数中查询，注意这个时候需要修改rank值

QUERYBYRANK(T, r)
    y = root[T]
    while y != nil
        if count[left[y]] + 1 == r then
            // find it
            exit
        else if count[left[y]] +1 > r then
            y = left[y]
        else
            r = r - count[left[y]] - 1
            y = right[y]

3.4 查询排名

红黑树普通查询，O(logn)可以查询到指定元素的排名

RANK(T, x)
    y = root[T]
    while y != nil
        if key[x] == key[y] then
            // find it
            r = count[y]
            exit
        else if key[x] < key[y] then
            y = left[y]
        else y = right[y]