可持久化线段树(主席树)

What · 是什么

可持久化线段树（Persistent Segment Tree，中文俗称『主席树』）是保存每次修改后历史版本的线段树。核心思想是 path copying：每次点修改只新建 O(log n) 个『被改路径』上的新节点，未变子树直接复用旧节点指针，因此 m 次修改的总节点数 O(m log n)。权值线段树版（每个叶子代表一个值域桶，维护区间内该值出现的次数）配合前缀和差分，可在 O(log n) 回答区间第 k 小/第 k 大、区间 ≤x 的个数、静态/动态区间众数类查询。经典应用：POJ 2104 K-th Number（区间第 k 小）、Luogu P3834/P2633 等。是竞赛中最常考的可持久化结构。

Why · 为什么

How · 实现

伪代码

ls[], rs[], cnt[], tot=0; rt[0]=build();
update(prev, l, r, pos):
  cur = ++tot; ls[cur]=ls[prev]; rs[cur]=rs[prev]; cnt[cur]=cnt[prev]+1
  if l<r:
    mid=(l+r)//2
    if pos<=mid: ls[cur]=update(ls[prev], l, mid, pos)
    else:        rs[cur]=update(rs[prev], mid+1, r, pos)
  return cur
query(u, v, l, r, k):   # 用 [1,l-1] 与 [1,r] 两根差分
  if l==r: return l
  left = cnt[ls[v]] - cnt[ls[u]]
  if k<=left: return query(ls[u], ls[v], l, mid, k)
  else:       return query(rs[u], rs[v], mid+1, r, k-left)

Python 实现

from bisect import bisect_left


class PersistentSegTree:
    """可持久化权值线段树（主席树）。构造 O(n log n)，查询第 k 小 O(log n)。
    arr: 原序列。支持 kth(l, r, k) 返回区间 [l,r]（1 索引含端点）第 k 小值。"""

    def __init__(self, arr):
        self.vals = sorted(set(arr))
        self.ls = [0]
        self.rs = [0]
        self.cnt = [0]
        self.roots = [self._build(0, len(self.vals) - 1)]
        for x in arr:
            self.roots.append(
                self._update(self.roots[-1], 0, len(self.vals) - 1,
                             bisect_left(self.vals, x)))

    def _build(self, l, r):
        idx = len(self.ls)
        self.ls.append(0); self.rs.append(0); self.cnt.append(0)
        if l < r:
            mid = (l + r) // 2
            self.ls[idx] = self._build(l, mid)
            self.rs[idx] = self._build(mid + 1, r)
        return idx

    def _update(self, prev, l, r, pos):
        idx = len(self.ls)
        self.ls.append(self.ls[prev])
        self.rs.append(self.rs[prev])
        self.cnt.append(self.cnt[prev] + 1)
        if l < r:
            mid = (l + r) // 2
            if pos <= mid:
                self.ls[idx] = self._update(self.ls[prev], l, mid, pos)
            else:
                self.rs[idx] = self._update(self.rs[prev], mid + 1, r, pos)
        return idx

    def kth(self, l, r, k):
        """区间 [l,r]（1 索引含端点）第 k 小值"""
        return self.vals[self._kth(self.roots[l - 1], self.roots[r],
                                  0, len(self.vals) - 1, k)]

    def _kth(self, u, v, l, r, k):
        if l == r:
            return l
        mid = (l + r) // 2
        left = self.cnt[self.ls[v]] - self.cnt[self.ls[u]]
        if k <= left:
            return self._kth(self.ls[u], self.ls[v], l, mid, k)
        return self._kth(self.rs[u], self.rs[v], mid + 1, r, k - left)


if __name__ == '__main__':
    import random

    def brute(a, l, r, k):
        return sorted(a[l - 1:r])[k - 1]

    a = [1, 5, 2, 6, 3, 7, 4]
    pst = PersistentSegTree(a)
    assert pst.kth(2, 5, 2) == 3      # [5,2,6,3] 排序 [2,3,5,6] 第2小=3
    assert pst.kth(1, 7, 4) == 4
    assert pst.kth(3, 7, 5) == 7

    random.seed(21)
    for _ in range(200):
        n = random.randint(1, 12)
        arr = [random.randint(-20, 20) for _ in range(n)]
        pt = PersistentSegTree(arr)
        for _ in range(20):
            l = random.randint(1, n)
            r = random.randint(l, n)
            k = random.randint(1, r - l + 1)
            assert pt.kth(l, r, k) == brute(arr, l, r, k), (l, r, k, arr)
    print('persistent-seg: All tests passed.')

易错点

• 动态开点：不能像普通线段树用 x*2/x*2+1 堆式存储，必须显式存 ls/rs
• update 必须先拷贝 prev 的全部字段（ls/rs/cnt）再改一个分支，漏拷贝导致历史被破坏
• 离散化后查询返回的是值域下标，输出前要映射回原值
• 空间：节点上限不足（< n*log n）会 RE/内存溢出，OI 惯例开 n<<5
• 区间 [l,r] 用 rt[r] 与 rt[l-1] 两根差分，用成 rt[l] 会答案错误
• 第 k 大 = 第 (len-k+1) 小，转换易混；k 必须 ≤ r-l+1 否则越界

例题

信源

• OI Wiki✓ https://oi-wiki.org/ds/persistent-seg/
• OI Wiki✓ https://raw.githubusercontent.com/OI-wiki/OI-wiki/master/docs/ds/persistent-seg.md
• 书籍✓ CPH ch28 pp271-273