稀有猿诉

十年磨一剑,历炼出锋芒,说话千百句,不如码二行。

线段树让你不再惧怕区间问题

线段树是用于求解序列区间问题的高效的数组结构,它是以完全二叉树形式来把序列划分为不同的区间,进而把O(n)复杂度提升到O(logn)。

线段树与树状数组有些类似,都可以用于解决区间问题,但线段树更具体普适性,树状数组能解决的问题线段树都可以解决,但反之不然。树状数组的学习可以参考另外一篇文章

首先,明确区间问题,是指对一个序列(或者数组)的某一个区间[left, right]做操作,比如区间求和,区间最值,区间修改(区间内加上或者减去某一数值),单点查询和单点修改(其实单点查询和单点修改可以视为区间查询区间修改的特殊情况即left=right时)。当遇到区间问题时,就可以考虑使用线段树来求解。

线段树的原理

线段树是一个辅助数据结构,以一个完全二叉树的形式把原数组(序列)展示出来,每一个节点代表着原序列的一个区间和区间查询。这里区间查询的意思是与区间相关的一个值,比如区间和,区间最值。区间顺着二叉树DFS向下不断二分变小,根节点是整个序列区间[0, n - 1],左子节点则是[0, mid],右子节点是[mid+1, n-1],其中mid=(n-1)/2。叶子节点就是原序列的单个元素,叶了节点即是长度为0的区间,即[left, right],且left=right。

比如,一个序列(数组)nums = {10, 11, 12, 13, 14},将其转化为线段树,区间查询 方式为区间求和,就会是这个样子的:

flowchart TD; t1(["tree[1]=60 [0, 4]"]) t2(["tree[2]=33 [0, 2]"]) t3(["tree[3]=27 [3, 4]"]) t4(["tree[4]=21 [0, 1]"]) t5(["tree[5]=12 [2, 2]"]) t6(["tree[6]=13 [3, 3]"]) t7(["tree[7]=14 [4, 4]"]) t8(["tree[8]=10 [0, 0]"]) t9(["tree[9]=11 [1, 1]"]) t1-->t2; t1-->t3; t2-->t4; t2-->t5; t4-->t8; t4-->t9; t3-->t6; t3-->t7;

从这个图中可以看出线段树的父节点与子节点的关系,tree[k] = tree[2 * k] + tree[2 * k + 1],也就是前面说的父节点的区间分成两半给左子节点和右子节点。

注意,线段树的编号从1开始,这是为了方便做DFS,当用数组作为具体实现时要注意tree的下标从1开始,这样的话父节点能非常方便的找到它的两个子节点(即2*i和2*i+1),而子节点也能方便的找到它的父节点(即k/2)。这样1是根,2和3是它的子节点,2和3除2都是1。但如果编号从0开始,就麻烦了,父子关系的运算不可逆,0找不到左子节点(当然 也不是完全不可以,需要做些额外的加1动作)。而且要注意这个下标与具体的区间没有任何关系,整个线段树的区间仍是输入数组的区间,即[0, n-1]

线段树的操作都是用从根开始的DFS来完成,因为这是一个平衡二叉树树的高度最高为logn,所以线段树的操作如单点更新和区间查询的复杂度都是O(logn)的。需要注意的是线段树的区间是通过节点数量以及节点之间的关系来体现的,节点存储的是区间查询(如区间和,最值),并没有具体存储区间信息。一切查询都从根节点开始,根节点包含整个序列区间,然后每往下走就区间减半,直到只包含一个序列的叶子节点。

核心原理就是,根节点包含序列的整个区间,不断向下DFS,这个过程中区间会减半收缩,当DFS过程中的区间小于等于目标区间时,就可以停止DFS了,因为已找到了答案,此时节点存储的查询值就是答案。

外部调用线段树进行查询时,必须提供目标区间作为参数(单点问题可视为特殊区间),这是因为从根往下做DFS目标区间是终止信号,在后面的实现细节时可以看到。因此,线段树只适合解决区间相关的问题,原因就在于此。

基础线段树的实现

先实现一个基础版本的线段树,学习建树的方法,以及支持单点修改区间查询(即区间求和)

需要注意的线段树,是一颗真实的二叉树,节点存储的是区间查询(区间和或者最值),节点之间的关系体现了区间。但在具体的实现上面,仍然可以用一个数组,严谨的说法是堆式存储(用二叉堆的方式),节点的关系就是父节点的索引是i,那么左子节点就在2*i,右子节点就在2*+1。因为线段对是平衡二叉树,所有操作又都从根节点开始,所以用堆式存储是比较方便的。

需要注意线段树数组的长度并不是原序列长度,因为节点数量取决 于区间的个数,所以是远大于原序列长度的,比如原数组长度是n,那么线段树长度约为4*n,具体的证明可以看后面列出的参考文章(其实我也看不太懂具体的证明)。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
/**
 * Segment Tree
 * Basic version, support point update (PU) and region query (RU).
 * Underlying is an array with 4*n size, where n is the length of input array.
 */
public class BasicSegmentTree {
    private static final int BASE = 1;

    private final int[] tree;

    /**
     * The length of input array.
     */
    private final int size;

    /**
     * Instantiate the segment tree.
     * @param nums the input array.
     */
    public BasicSegmentTree(int[] nums) {
        size = nums.length;

        tree = new int[4 * size];

        build(nums, 0, size - 1, BASE);
    }

    /**
     * Replace item at k with val.
     * @param k the index to replace in input array.
     * @param val target value to replace with.
     */
    public void update(int k, int val) {
        doUpdate(k, val, 0, size - 1, BASE);
    }

    /**
     * Get the region sum between [left, right] inclusive.
     * @param left the left index of the region in input array, inclusive.
     * @param right the right index of the region in input array, inclusive.
     * @return the region sum.
     */
    public int query(int left, int right) {
        return doQuery(left, right, 0, size - 1, BASE);
    }

    private void build(int[] nums, int start, int end, int index) {
        if (start == end) {
            tree[index] = nums[start];
            return;
        }
        int mid = start + ((end - start) >> 1);
        build(nums, start, mid, 2 * index);
        build(nums, mid + 1, end, 2 * index + 1);
        pushUp(index);
    }

    private void pushUp(int index) {
        tree[index] = tree[2 * index] + tree[2 * index + 1];
    }

    private void doUpdate(int k, int val, int start, int end, int index) {
        if (start == end) {
            tree[index] = val;
            return;
        }
        int mid = start + ((end - start) >> 1);
        if (k <= mid) {
            doUpdate(k, val, start, mid, 2 * index);
        } else {
            doUpdate(k, val, mid + 1, end, 2 * index + 1);
        }
        pushUp(index);
    }

    /**
     * Get the sum of target region [left, right] starting from current region [start, end].
     * @param left target region left index, inclusive.
     * @param right target region right index, inclusive.
     * @param start current region left index, inclusive.
     * @param end current region right index, inclusive.
     * @param index current index of the segment tree.
     * @return the sum of target region.
     */
    private int doQuery(int left, int right, int start, int end, int index) {
        if (left <= start && end <= right) {
            return tree[index];
        }
        int sum = 0;
        int mid = start + ((end - start) >> 1);
        /*
         * Case #1: only need to query left half of current region.
         *         [left, right]
         *                       mid
         * Case #2: only need to query right half
         *                            [left, right]
         *                       mid
         * Case #3: query both left and right half
         *                [left,               right]
         *         [start         mid, mid+1           end]
         * Limitation: left <= right, start <= end
         * So, if left > mid, must have mid < left <= right, which is Case #2;
         * if right <= mid, must have left <= right <= mid, which is Case #1;
         * if left <= mid < right, which is Case #3.
         */
        if (left <= mid) {
            sum += doQuery(left, right, start, mid, 2 * index);
        }
        if (right > mid) {
            sum += doQuery(left, right, mid + 1, end, 2 * index + 1);
        }
        return sum;
    }

    public static void main(String[] args) {
        int[] a = {10, 11, 12, 13, 14};

        BasicSegmentTree bst = new BasicSegmentTree(a);

        System.out.println("60 = " + bst.query(0, 4));
        System.out.println("33 = " + bst.query(0, 2));
        System.out.println("36 = " + bst.query(1, 3));
        System.out.println("27 = " + bst.query(3, 4));

        bst.update(2, 15);
        System.out.println("63 = " + bst.query(0, 4));
        System.out.println("36 = " + bst.query(0, 2));
        System.out.println("39 = " + bst.query(1, 3));
        System.out.println("27 = " + bst.query(3, 4));
    }
}

完整代码参见这里

单点查询的支持

前面实现了最基础的区间操作:单点修改(PU)和区间查询(RQ),下面来看一下单点查询(PQ)如何实现。其实单点查询特别简单,把单点更新稍改一下就是单点查询了:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
/**
 * Query the item at index of the input array.
 * @param index of the input array.
 * @return the item at index.
 */
public int pointQuery(int index) {
    return doPointQuery(index, 0, size - 1, BASE);
}

private int doPointQuery(int where, int start, int end, int index) {
    if (start == end) {
        return tree[index];
    }
    int mid = start + ((end - start) >> 1);
    if (where <= mid) {
        return doPointQuery(where, start, mid, 2 * index);
    } else {
        return doPointQuery(where, mid + 1, end, 2 * index + 1);
    }
}

当然,这里可以取巧,因为单点查询可以视为区间查询的特殊情况,因此如果直接调用doQuery(index, index, 0, size - 1, 1)来实现pointQuery()也是可行的。但毕竟doQuery中做了累加的动作,会有额外开销,单独实现一下要好一些。

还有,如果在实现中保存了输入数组,对于单点查询 就更好办了,直接返回nums[index]就好了。但有一个问题,就是保存了输入数组,就需要维护它的状态,因为不光有查询 ,还有更新,要在单点更新时也同时更新一下输入数组。

区间更新

另外一个区间问题就是区间更新,也就是对于给定的区间内每个元素都加上一个值。

虽然,线段树节点中会存储区间查询(比如区间和),看似乎DFS时我们找到了目标区间后,把此节点的查询值改一下就可以了,比如前面的例子,add(3, 4, 10),当找到目标区间时,即根节点的右子节点,是tree[3],把tree[3] 加上 2 * 10(区间内元素个数 为2,节点存储的是区间和,所以要加上个数乘上数值)。

看似是对的,但子区间的信息并没有更新啊,假如一个查询刚好是[3, 4],你能得到正确的结果,但如果查询[2, 4]或者单点查询3或者单点查询4,因为这些节点并未更新,还是原来的值,查询结果就会是错误的。

你可能的会说,那就顺带更新一下子区间(即子节点)不就完事儿了么,比如在add(3, 4, 10),更新了tree[3]后,再把它的子节点tree[6]和tree[7]也更新一下,后面再怎么查询都能得到正确的结果了。

确实可以这样做,但时间复杂度会上升到O(n),并且因为线段树的节点数量是大于输入数组的,所以实际的效率甚至比更新输入数组(纯O(n))的效率还要差。

为了解决问题,就需要用到叫做『懒惰标记』的技巧。新开一个与tree等大的数组lazy专门用于标记哪些区间被更新了,然后当有做向下DFS的机会时,进行『标记下放』,这样就可以保证更新只做一次。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
   /**
     * Add x to all items between [left, right] inclusive.
     * @param left left index of the region.
     * @param right right index of the region.
     * @param x value to add.
     */
    public void add(int left, int right, int x) {
        if (x == 0) {
            return;
        }
        doAdd(left, right, x, 0, size - 1, BASE);
    }

    private void doAdd(int left, int right, int x, int start, int end, int index) {
        if (left <= start && end <= right) {
            tree[index] += (end - start + 1) * x;
            if (start != end) {
                lazy[index] += x;
            }
            return;
        }

        int mid = start + ((end - start) >> 1);

        if (lazy[index] != 0) {
            pushDown(start, mid, end, index);
        }

        if (left <= mid) {
            doAdd(left, right, x, start, mid, 2 * index);
        }
        if (right > mid) {
            doAdd(left, right, x, mid + 1, end, 2 * index + 1);
        }
        pushUp(index);
    }

    private void pushDown(int start, int mid, int end, int index) {
        tree[2 * index] += (mid - start + 1) * lazy[index];
        lazy[2 * index] += lazy[index];

        tree[2 * index + 1] += (end - mid) * lazy[index];
        lazy[2 * index + 1] += lazy[index];

        lazy[index] = 0;
    }

具体的逻辑就是,当做区间更新时,找到了目标区间,先给目标区间的节点做更新,然后标记这个节点有过更新,此是『懒惰标记』,更新过程到此结束。之后,在所有的DFS过程中(任何的更新和查询),都会看当前区间,是否有被标记过,如果有被标记过,那么就要更新一下它的两个子区间,此称之为『标记下放』。注意每次『标记下放』只是从当前区间下放到其子区间,并没有直接DFS下去,什么时候当访问子区间时,再去继续向下下放,这就体现了懒惰的意思。这样就能保证,区间更新的效果只向其子区间遍历一次,并且在任何一次DFS时就顺带下去了,不会产生额外的复杂度开销,是O(1)的。综合分析,区间更新的复杂度只消耗在更新目标区间的过程,即是O(logn)的。通过『懒惰标记』方法,把子区间的更新摊还到了其他DFS过程中,不产生额外的开销,所以总的复杂度是O(logn)。

完整代码参见这里

区间替换

还有另外一种区间更新方式,就是把区间内所有元素都替换为另外一个值。整体的思路与区间更新是一样的,就是修改节点时改成赋值而不是累加。但在做懒惰标记需要注意,因为新值可能为0,所以判断一个节点是否有标记过,是否需要『标记下放』时,就不能用lazy[index] == 0来判断了,视题目数据 范围,可以用一个常规数值范围外的值来当成默认未标记的状态,或者再创建一个数组visited来用以查询节点是否标记过:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
   /**
     * Replace all items between [left, right] inclusive with val.
     * @param left left index of the region.
     * @param right right index of the region.
     * @param val value to replace
     */
    public void update(int left, int right, int val) {
        doUpdate(left, right, val, 0, size - 1, BASE);
    }

    private void doUpdate(int left, int right, int val, int start, int end, int index) {
        if (left <= start && end <= right) {
            tree[index] = (end - start + 1) * val;
            if (start != end) {
                lazy[index] = val;
                visited[index] = true;
            }
            return;
        }

        int mid = start + ((end - start) >> 1);

        if (visited[index]) {
            pushDown(start, mid, end, index);
        }

        if (left <= mid) {
            doAdd(left, right, val, start, mid, 2 * index);
        }
        if (right > mid) {
            doAdd(left, right, val, mid + 1, end, 2 * index + 1);
        }
        pushUp(index);
    }

    private void pushDown(int start, int mid, int end, int index) {
        tree[2 * index] = (mid - start + 1) * lazy[index];
        lazy[2 * index] = lazy[index];
        visited[2 * index] = true;

        tree[2 * index + 1] = (end - mid) * lazy[index];
        lazy[2 * index + 1] = lazy[index];
        visited[2 * index + 1] = true;

        lazy[index] = 0;
        visited[index] = false;
    }

动态开点

从前面的实现中可以看出线段树的空间复杂度是较高的,如果用堆式存储(即数组)树本身就要4倍的空间,再加上标记可能会提升一个数量级,在某些情况下,这是不可取的,比如输入数组特别长,或者输入数据产生的序列长度特别大,但有效元素非常少,会MLE。

这时就可以用动态开点,也就是动态的创建子节点,当DFS到某个子节点时,发现还没有再去创建。可以参考参考 资料中的文章。

指针节点

线段树是一个真正的二叉树,所以可以用常规指针节点的方式来构建树,这样动态创建每个节点,每个节点的字段也可以扩展,其实在某些情况下更为方便。具体实现方式也可以参考后面的文章。

1
2
3
4
5
6
7
8
9
10
class SegmentTree {
  Node root;
  
  class Node {
      int sum;
      int lazy;
      Node left;
      Node right;
  }
}

典型问题

题目 题解 说明
307. 区域和检索 - 数组可修改 题解 单点修改,区间查询,基础线段树
1450. 在既定时间做作业的学生人数 题解
731. 我的日程安排表 II 题解 动态开点,区改,区间最大值
729. 我的日程安排表 I 题解 动态开点,区改区查
56. 合并区间 题解
题解
题解

参考资料

Comments