线段树

2024-10-05

字数统计: 3.6k字 | 阅读时长: 16分

前言

线段树是算法竞赛中常用的用来维护“区间信息”的数据结构。线段树可以在 $O(\log N)$ 的时间复杂度内实现单点修改、区间修改、区间查询（区间求和，求区间最大值，求区间最小值）等操作。

线段树（Segment Tree）

线段树将每个长度不为$1$的区间划分成左右两个区间递归求解，把整个线段划分为一个树形结构，通过合并左右两区间信息来求得该区间的信息。这种数据结构可以方便的进行大部分的区间操作。

基本结构

有个大小为$5$的数组 a={10,11,12,13,14}，要将其转化为线段树，有以下做法：设线段树的根节点编号为$1$，用数组 $d$ 来保存我们的线段树，$d_i$ 用来保存线段树上编号为 $i$ 的节点的值（这里每个节点所维护的值就是这个节点所表示的区间总和）。这棵线段树的形态如下图所示：

segt1

图中每个节点中用红色字体标明的区间，表示该节点管辖的 $a$ 数组上的位置区间。如 $d_1$ 所管辖的区间就是 $[1,5]$（$a_1,a_2, \cdots ,a_5$），即 $d_1$ 所保存的值是 $a_1+a_2+ \cdots +a_5$，$d_1=60$ 表示的是 $a_1+a_2+ \cdots +a_5=60$。

通过观察不难发现，$d_i$ 的左儿子节点就是 $d_{2\times i}$，$d_i$ 的右儿子节点就是 $d_{2\times i+1}$。如果 $d_i$ 表示的是区间 $[s,t]$（即 $d_i=a_s+a_{s+1}+ \cdots +a_t$）的话，那么 $d_i$ 的左儿子节点表示的是区间 $[ s, \frac{s+t}{2} ]$，$d_i$ 的右儿子表示的是区间 $[ \frac{s+t}{2} +1,t ]$。

建树过程

可以使用递归建树，设当前的根节点为 $p$，如果根节点管辖的区间长度已经是 $1$，则可以直接根据 $a$ 数组上相应位置的值初始化该节点。否则我们将该区间从中点处分割为两个子区间，分别进入左右子节点递归建树，最后合并两个子节点的信息。

采用数组堆式存储（$2p$ 是 $p$ 的左儿子，$2p+1$ 是 $p$ 的右儿子），若有 $n$ 个叶子结点，则数组的范围最大为 $2^{\left\lceil\log{n}\right\rceil+1}$。

分析：容易知道线段树的深度是 $\left\lceil\log{n}\right\rceil$ 的，则在堆式储存情况下叶子节点（包括无用的叶子节点）数量为 $2^{\left\lceil\log{n}\right\rceil}$ 个，又由于其为一棵完全二叉树，则其总节点个数 $2^{\left\lceil\log{n}\right\rceil+1}-1$。当然也可以直接把数组长度设为 $4n$，因为 $\frac{2^{\left\lceil\log{n}\right\rceil+1}-1}{n}$ 的最大值在 $n=2^{x}+1(x\in N_{+})$ 时取到，此时节点数为 $2^{\left\lceil\log{n}\right\rceil+1}-1=2^{x+2}-1=4n-5$。

区间查询

仍然以最开始的图为例，如果要查询区间 $[1,5]$ 的和，那直接获取 $d_1$ 的值即可。

如果要查询的区间为 $[3,5]$，此时就不能直接获取区间的值，但是 $[3,5]$ 可以拆成 $[3,3]$ 和 $[4,5]$，可以通过合并这两个区间的答案来求得这个区间的答案。

一般地，如果要查询的区间是 $[l,r]$，则可以将其拆成最多为 $O(\log n)$ 个“极大”的区间，合并这些区间即可求出 $[l,r]$ 的答案。

区间修改与懒惰标记

如果要求修改区间 $[l,r]$，把所有包含在区间 $[l,r]$ 中的节点都遍历一次、修改一次，时间复杂度无法承受。我们这里要引入一个叫做「懒惰标记」的东西。

懒惰标记，简单来说，就是通过延迟对节点信息的更改，从而减少可能不必要的操作次数。每次执行修改时，我们通过打标记的方法表明该节点对应的区间在某一次操作中被更改，但不更新该节点的子节点的信息。实质性的修改则在下一次访问带有标记的节点时才进行。

Java线段树模版（4n内存）

需要注意的是，Java的泛型由于存在擦除，因此在编译后实际上并不能拿到类型信息，所有的类型都是基类，导致不能实现算数运算，在限定Number的情况下，可以使用double值计算结果，需要在外部调用后进行类型转换。

可以区间加/求和的线段树模板：

// author: SilenceZheng66
import java.util.Arrays;

public class SegTreeLazyRangeAdd<T extends Number> {
    private double[] tree; // 使用 double 数组来存储树节点
    private double[] lazy; // 懒惰标记数组
    private T[] arr; // 原始数组
    private int n, root, n4, end;

    public SegTreeLazyRangeAdd(T[] v) {
        n = v.length;
        n4 = n * 4;
        tree = new double[n4];
        lazy = new double[n4];
        arr = v;
        Arrays.fill(tree, 0);
        Arrays.fill(lazy, 0);
        end = n - 1;
        root = 1;
        build(0, end, 1);
        arr = null; // Clear reference to the original array
    }

    private void maintain(int cl, int cr, int p) {
        int cm = cl + (cr - cl) / 2;
        if (cl != cr && lazy[p] != 0) {
            lazy[p * 2] += lazy[p]; // 更新左子树的懒惰标记
            lazy[p * 2 + 1] += lazy[p]; // 更新右子树的懒惰标记
            tree[p * 2] += lazy[p] * (cm - cl + 1); // 更新左子树的值
            tree[p * 2 + 1] += lazy[p] * (cr - cm); // 更新右子树的值
            lazy[p] = 0; // 重置当前节点的懒惰标记
        }
    }

    private double rangeSum(int l, int r, int cl, int cr, int p) {
        if (l <= cl && cr <= r) return tree[p]; // 当前子区间被查询区间完全覆盖
        int m = cl + (cr - cl) / 2;
        double sum = 0; // 使用 double 类型进行求和
        maintain(cl, cr, p); // 维护当前节点
        if (l <= m) sum += rangeSum(l, r, cl, m, p * 2); // 左子树
        if (r > m) sum += rangeSum(l, r, m + 1, cr, p * 2 + 1); // 右子树
        return sum;
    }

    private void rangeAdd(int l, int r, double val, int cl, int cr, int p) {
        if (l <= cl && cr <= r) { // 如果当前区间是要更新区间的子区间
            lazy[p] += val; // 懒惰标记增加
            tree[p] += (cr - cl + 1) * val; // 更新当前节点的值
            return;
        }
        int m = cl + (cr - cl) / 2;
        maintain(cl, cr, p); // 维护当前节点
        if (l <= m) rangeAdd(l, r, val, cl, m, p * 2); // 左子树
        if (r > m) rangeAdd(l, r, val, m + 1, cr, p * 2 + 1); // 右子树
        tree[p] = tree[p * 2] + tree[p * 2 + 1]; // 更新当前节点的值
    }

    private void build(int s, int t, int p) {
        if (s == t) {
            tree[p] = arr[s].doubleValue(); // 使用 double 类型
            return;
        }
        int m = s + (t - s) / 2;
        build(s, m, p * 2); // 构建左子树
        build(m + 1, t, p * 2 + 1); // 构建右子树
        tree[p] = tree[p * 2] + tree[p * 2 + 1]; // 更新当前节点的值
    }

    public double rangeSum(int l, int r) {
        return rangeSum(l, r, 0, end, root);
    }

    public void rangeAdd(int l, int r, double val) {
        rangeAdd(l, r, val, 0, end, root);
    }
}

可以区间修改/求和的线段树模板:

// author: SilenceZheng66
import java.util.Arrays;

public class SegTreeLazyRangeSet {
    private double[] tree; // 存储树节点的数组
    private double[] lazy; // 懒惰标记数组
    private boolean[] ifLazy; // 懒惰标记的状态
    private double[] arr; // 原始数组
    private int n, root, n4, end;

    public SegTreeLazyRangeSet(double[] v) {
        n = v.length;
        n4 = n * 4;
        tree = new double[n4];
        lazy = new double[n4];
        ifLazy = new boolean[n4];
        arr = v;
        Arrays.fill(tree, 0);
        Arrays.fill(lazy, 0);
        Arrays.fill(ifLazy, false);
        end = n - 1;
        root = 1;
        build(0, end, 1);
        arr = null;
    }

    private void maintain(int cl, int cr, int p) {
        int cm = cl + (cr - cl) / 2;
        if (cl != cr && ifLazy[p]) {
            lazy[p * 2] = lazy[p];
            ifLazy[p * 2] = true;
            lazy[p * 2 + 1] = lazy[p];
            ifLazy[p * 2 + 1] = true;
            tree[p * 2] = lazy[p] * (cm - cl + 1);
            tree[p * 2 + 1] = lazy[p] * (cr - cm);
            lazy[p] = 0;
            ifLazy[p] = false; // 重置当前节点的懒惰标记
        }
    }

    private double rangeSum(int l, int r, int cl, int cr, int p) {
        if (l <= cl && cr <= r) return tree[p]; // 完全覆盖
        int m = cl + (cr - cl) / 2;
        double sum = 0; // 使用 double 类型进行求和
        maintain(cl, cr, p); // 维护当前节点
        if (l <= m) sum += rangeSum(l, r, cl, m, p * 2); // 左子树
        if (r > m) sum += rangeSum(l, r, m + 1, cr, p * 2 + 1); // 右子树
        return sum;
    }

    private void rangeSet(int l, int r, double val, int cl, int cr, int p) {
        if (l <= cl && cr <= r) {
            lazy[p] = val; // 懒惰标记增加
            ifLazy[p] = true; // 标记为懒惰
            tree[p] = (cr - cl + 1) * val; // 更新当前节点的值
            return;
        }
        int m = cl + (cr - cl) / 2;
        maintain(cl, cr, p); // 维护当前节点
        if (l <= m) rangeSet(l, r, val, cl, m, p * 2); // 左子树
        if (r > m) rangeSet(l, r, val, m + 1, cr, p * 2 + 1); // 右子树
        tree[p] = tree[p * 2] + tree[p * 2 + 1]; // 更新当前节点的值
    }

    private void build(int s, int t, int p) {
        if (s == t) {
            tree[p] = arr[s]; // 使用 double 类型
            return;
        }
        int m = s + (t - s) / 2;
        build(s, m, p * 2); // 构建左子树
        build(m + 1, t, p * 2 + 1); // 构建右子树
        tree[p] = tree[p * 2] + tree[p * 2 + 1]; // 更新当前节点的值
    }

    public double rangeSum(int l, int r) {
        return rangeSum(l, r, 0, end, root);
    }

    public void rangeSet(int l, int r, double val) {
        rangeSet(l, r, val, 0, end, root);
    }
}

线段树优化

在叶子节点处无需下放懒惰标记，所以懒惰标记可以不下传到叶子节点。
下放懒惰标记可以写一个专门的函数maintain，降低代码编写难度。
标记永久化：如果确定懒惰标记不会在中途被加到溢出（即超过了该类型数据所能表示的最大范围），那么就可以将标记永久化。标记永久化可以避免下传懒惰标记，只需在进行询问时把标记的影响加到答案当中，从而降低程序常数。具体如何处理与题目特性相关，需结合题目来写。这也是树套树和可持久化数据结构中会用到的一种技巧。
动态开点线段树，即可以不一次性建好树，而是在最初只建立一个根结点代表整个区间。当我们需要访问某个子区间时，才建立代表这个区间的子结点。这种方式节省的空间有限（约$1/4$）。

其他操作

线段树的其他操作包括线段树合并、线段树分裂、线段树二分、线段树优化建图等，遇到相关题目时再做补充。

例题1: Booking Concert Tickets in Groups

将每一排座位看成一个整体，可将两种操作总结如下：

gather：在前 maxRow 排中，找第一个还能至少有 k 个空座位的排，预定 k 个连续的座位。如果有这样的排，返回编号，以及在预定前有多少非空座位；如果没有这样的排，返回空列表。
scatter：在前 maxRow 个排中预定总量为 k 的座位。从左到右选择有空座位的排依次占座。如果无法预定总量为 k 的座位，则不执行操作，并返回 false；否则执行操作，并返回 true。

此时将原问题压缩为一维数组，能这样做的原因是两种操作均不会产生“间隔的空座位”。因此需要的操作如下：

求出前 maxRow 排中，第一个剩余容量 ≥k，也就是已预定量 ≤ m−k 的排。
维护每排的已预定量。
维护前 maxRow 排的已预定量之和，从而判断 scatter 能否实现。

可以用线段树处理上述操作，线段树维护每个区间（区间即某几排）已预定量的最小值 $\min$ ，以及每个区间的已预定量之和 $sum$。

对于 gather，从线段树的根节点开始递归：

如果当前区间 $\min > m-k$, 则无法预定 $k$ 个连续座位，返回 0 。
如果当前区间长度为 1 ，返回区间端点。
如果左半区间 $\min \leq m-k$, 则答案在左半区间中, 递归左半区间。
否则如果 maxRow 在右半区间内, 递归右半区间。
否则返回 -1 ，表示没有符合条件的排。

上述过程叫做线段树二分，

对于 scatter，如果区间 $[0$, maxRow] 的已预定量之和大于 $m \cdot(\operatorname{maxRow}+1)-k$ ，则无法执行操作。否则可以执行操作。从第一个没有装满，也就是已预定量 $\leq m-1$ 的排开始预定，这也可以用线段树二分求出。

class BookMyShow {
    private int n;
    private int m;
    private int[] min;
    private long[] sum;

    public BookMyShow(int n, int m) {
        this.n = n;
        this.m = m;
        int size = 2 << (32 - Integer.numberOfLeadingZeros(n)); // 比 4n 更小
        min = new int[size];
        sum = new long[size];
    }

    public int[] gather(int k, int maxRow) {
        // 找第一个能预定 k 个座位的row
        int r = findFirst(1, 0, n - 1, maxRow, m - k);
        if (r < 0) { // 没有这样的row
            return new int[]{};
        }
        int c = (int) querySum(1, 0, n - 1, r, r);
        update(1, 0, n - 1, r, k); // 预定
        return new int[]{r, c};
    }

    public boolean scatter(int k, int maxRow) {
        // [0,maxRow] 的已预定量之和
        long s = querySum(1, 0, n - 1, 0, maxRow);
        if (s > (long) m * (maxRow + 1) - k) {
            return false; // 空座位不足
        }
        // 从第一个没有定满的row开始
        int i = findFirst(1, 0, n - 1, maxRow, m - 1);
        while (k > 0) {
            int left = Math.min(m - (int) querySum(1, 0, n - 1, i, i), k);
            update(1, 0, n - 1, i, left); // 预定
            k -= left;
            i++;
        }
        return true;
    }

    // 把下标 i 上的元素值增加 val
    private void update(int o, int l, int r, int i, int val) {
        if (l == r) {
            min[o] += val;
            sum[o] += val;
            return;
        }
        int m = (l + r) / 2;
        if (i <= m) {
            update(o * 2, l, m, i, val);
        } else {
            update(o * 2 + 1, m + 1, r, i, val);
        }
        min[o] = Math.min(min[o * 2], min[o * 2 + 1]);
        sum[o] = sum[o * 2] + sum[o * 2 + 1];
    }

    // 返回区间 [L,R] 内的元素和
    private long querySum(int o, int l, int r, int L, int R) {
        if (L <= l && r <= R) {
            return sum[o];
        }
        long res = 0;
        int m = (l + r) / 2;
        if (L <= m) {
            res = querySum(o * 2, l, m, L, R);
        }
        if (R > m) {
            res += querySum(o * 2 + 1, m + 1, r, L, R);
        }
        return res;
    }

    // 返回区间 [0,R] 中 <= val 的最靠左的位置，不存在时返回 -1
    private int findFirst(int o, int l, int r, int R, int val) {
        if (min[o] > val) {
            return -1; // 整个区间的元素值都大于 val
        }
        if (l == r) {
            return l;
        }
        int m = (l + r) / 2;
        if (min[o * 2] <= val) {
            return findFirst(o * 2, l, m, R, val);
        }
        if (R > m) {
            return findFirst(o * 2 + 1, m + 1, r, R, val);
        }
        return -1;
    }
}

参考文献

[1] https://oi-wiki.org/ds/seg/#%E5%AE%9E%E7%8E%B0
[2] https://en.wikipedia.org/wiki/Segment_tree
[3] https://leetcode.cn/problems/booking-concert-tickets-in-groups/solutions/1523876/by-endlesscheng-okcu/?envType=daily-question&envId=2024-09-28

后记

首发于 silencezheng.top，转载请注明出处。