【模板】普通平衡树

GitHub跳转原题关系图返回列表

离线收集所有数值并坐标压缩,用 Fenwick 维护多重集合、排名、kth、前驱和后继。

OJ: luogu

题目 ID: P3369

难度:提高+/省选-

标签:树状数组坐标压缩有序多重集python

日期: 2026-07-16 19:57

题意

动态维护可重集合,支持插入、删除、排名、第 k 小、严格前驱与严格后继。

思路

题目虽然叫平衡树,但全部操作已经在输入中,可以先离线收集所有作为“数值”的 x 并坐标压缩。树状数组在压缩坐标上维护每个值的出现次数。

  • 排名:小于 x 的数量 + 1
  • k 小:在 Fenwick 树上二进制提升,找最小前缀和达到 k 的坐标;
  • 前驱:小于 x 的数量所对应的第 count 小;
  • 后继:小于等于 x 的数量再加一所对应的元素。

重复值只改变同一坐标的计数,语义与可重集合一致。

Python 知识

  • sorted({value for ...}) 一行完成离线去重排序。
  • bisect_left 统计严格小于,bisect_right 统计小于等于。
  • Fenwick 的 kth 使用 bit_length() 取得最高二进制步长。
  • 操作保存成元组列表,第二遍执行时无需重新解析。

代码

python
import sys
from bisect import bisect_left, bisect_right


data = iter(map(int, sys.stdin.buffer.read().split()))
operations = [(next(data), next(data)) for _ in range(next(data))]
coordinates = sorted({value for operation, value in operations if operation != 4})
tree = [0] * (len(coordinates) + 1)


def add(index, delta):
    index += 1
    while index < len(tree):
        tree[index] += delta
        index += index & -index


def prefix(end):
    result = 0
    while end:
        result += tree[end]
        end -= end & -end
    return result


def kth(rank):
    index = 0
    step = 1 << (len(tree).bit_length() - 1)
    while step:
        target = index + step
        if target < len(tree) and tree[target] < rank:
            index = target
            rank -= tree[target]
        step >>= 1
    return coordinates[index]


answers = []
for operation, value in operations:
    if operation == 1:
        add(bisect_left(coordinates, value), 1)
    elif operation == 2:
        add(bisect_left(coordinates, value), -1)
    elif operation == 3:
        answers.append(str(prefix(bisect_left(coordinates, value)) + 1))
    elif operation == 4:
        answers.append(str(kth(value)))
    elif operation == 5:
        answers.append(str(kth(prefix(bisect_left(coordinates, value)))))
    else:
        answers.append(str(kth(prefix(bisect_right(coordinates, value)) + 1)))

print("\n".join(answers))

复杂度

设操作数为 nn,预排序 O(nlogn)O(n\log n),每次操作 O(logn)O(\log n),空间 O(n)O(n)

总结

OJ 输入允许离线时,不必强行在 Python 手写旋转平衡树;坐标压缩 + Fenwick 同样完整实现顺序统计接口。