使用fenwick树(BIT)的订单统计树

简介

订单统计树(Order Statistic Tree，OST)是一种数据结构，支持在log(n)时间复杂度内，对于一些与二叉查找树有关的查询（例如：查找第k大元素）。

该数据结构可以使用红黑树作为底层实现，但是维护起来比较麻烦，而且常数比较大。另外，红黑树是不能实现动态数组的，而OST可以。

为了解决这些问题，可以使用fenwick树（也称为二叉索引树或树状数组）作为底层实现。fenwick树有很多用途，其中之一就是实现动态数组的查找和更新操作。

fenwick树的实现

一棵有n个元素的fenwick树可以用一个长度为n的数组表示，其中第i个元素存放了以i为最右边界时，前面所有元素的和。

对于任意一个下标i，其父结点的下标可以通过将i的二进制表示中最右边的1变成0来计算：

def parent(i):
    return i - (i & -i)

同样地，它的下一个包含在区间内的元素下标j可以通过将i的二进制表示中最右边的0变成1来计算：

def next(j):
    return j + (j & -j)

而我们要查询区间和时，可以使用“树状数组”的思想，从右往左把每个结点对应的区间和算出来，然后递推得到所求区间的和。

def query(i):
    res = 0
    while i > 0:
        res += tree[i]
        i = parent(i)
    return res

最后是修改操作：

def update(i, val):
    while i <= n:
        tree[i] += val
        i = next(i)

OST的实现

经过上面的介绍，我们可以将OST实现为一个动态维护升序序列的数据结构。对于任意一个节点，我们可以记录下它子树中的元素个数，从而将查询第k大元素的操作转化为树上的查找操作。

对于插入和删除操作，我们可以利用“树状数组”的思想，将所有受影响的节点的子树元素个数更新。具体实现方式与上面的fenwick树类似。

代码实现

下面是利用fenwick树实现的OST的代码。节点的结构体中包含了它的元素值value，它所在的下标index，以及它的子树元素个数size。

class Node:
    def __init__(self, val, idx):
        self.value = val
        self.index = idx
        self.size = 1

class FenwickTree:
    def __init__(self, n):
        self.n = n
        self.tree = [0 for _ in range(n+1)]

    def parent(self, i):
        return i - (i & -i)

    def next(self, j):
        return j + (j & -j)

    def query(self, i):
        res = 0
        while i > 0:
            res += self.tree[i]
            i = self.parent(i)
        return res

    def update(self, i, val):
        while i <= self.n:
            self.tree[i] += val
            i = self.next(i)

class OrderStatisticTree:
    def __init__(self):
        self.fenwick = FenwickTree(100000)
        self.root = None

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

    def kth(self, k):
        node = self.root
        while node:
            left_size = self.size(node.left)
            if left_size + 1 == k:
                return node.value
            elif left_size >= k:
                node = node.left
            else:
                k -= left_size + 1
                node = node.right

    def insert(self, val, idx):
        node = Node(val, idx)
        if not self.root:
            self.root = node
        else:
            parent = None
            cur = self.root
            while cur:
                cur.size += 1
                parent = cur
                if val < cur.value:
                    cur = cur.left
                else:
                    cur = cur.right
            if val < parent.value:
                parent.left = node
            else:
                parent.right = node
            node.size = 1
            self.fenwick.update(idx, 1)

    def delete(self, val, idx):
        parent = None
        node = self.root
        while node:
            if val == node.value:
                break
            parent = node
            if val < node.value:
                node = node.left
            else:
                node = node.right
        if not node:
            return
        self.fenwick.update(idx, -1)
        size = node.size
        if node.left and node.right:
            parent = node
            succ = node.right
            while succ.left:
                parent = succ
                succ = succ.left
            node.value = succ.value
            node.index = succ.index
            node = succ
        child = node.left if node.left else node.right
        if not parent:
            self.root = child
        elif parent.left == node:
            parent.left = child
        else:
            parent.right = child
        while parent:
            parent.size -= size
            parent = parent.parent

结语

fenwick树和OST都是非常有用的数据结构，它们的应用场景也非常广泛。熟练掌握它们的实现方法，可以在编程竞赛和工程实践中提高我们的效率和质量。