【小航的算法日记】线段树

本内容取经于：https://leetcode.cn/problems/my-calendar-i/solution/by-lfool-xvpv/

一、概念

概念区分：

• 线段树解决的是「区间和」的问题，且该「区间」会被修改
• 前缀和解决的是 「区间和」的问题，且该「区间」不会被修改

举个栗子：

对于nums = [1, 2, 3, 4, 5]，如果 nums 会被修改呢？比如：

• 把第 i 个元素修改成 x
• 把第 i 个元素增加 x
• 把区间 [i, j] 内的元素都增加 x

这个时候，如果我们使用前缀和，就没有那么高效了， 每一次更新前缀和数组都需要更新，时间复杂度为O(N)

---

原理及实现：

首先始终记住一句话：线段树的每个节点代表一个区间
由于线段树解决的是「区间和」的问题，且该「区间」会被修改，所以它主要实现两个方法：「求区间和」&&「修改区间」，且时间复杂度均为O(logn)

nums = [1, 2, 3, 4, 5] 对应的线段树如下所示：

从图中可以看到，每个节点代表一个区间，而节点的值就是该区间的和

当然节点的含义还可以是其他：

• 数字之和「总数字之和 = 左区间数字之和 + 右区间数字之和」
• 最大公因数 (GCD)「总 GCD = gcd(左区间 GCD, 右区间 GCD)」
• 最大值「总最大值 = max(左区间最大值，右区间最大值)」

不合法的栗子：

• 众数「只知道左右区间的众数，没法求总区间的众数」
• 01 序列的最长连续零「只知道左右区间的最长连续零，没法知道总的最长连续零」

根节点代表的区间是问题的总区间，如这个例子，问题的总区间就是数组的长度 [0, 4]

二、模板

线段树的数据结构:

class Node {
    // 左右***节点
    Node left, right;
    // 当前节点值
    int val;
}

我们使用数组来表示一棵线段树，假如根节点为 i，那么左的节点就为 2 * i，右的节点就为 2 * i + 1 (前提：i 从 1 开始)

线段树的建立：

题目给了具体的区间范围：

解法：根据该范围建立线段树

// 后序遍历
public void buildTree(Node node, int start, int end) {
    // 到达叶子节点
    if (start == end) {
        node.val = arr[start];
        return ;
    }
    int mid = (start + end) >> 1;
    buildTree(node.left, start, mid);
    buildTree(node.right, mid + 1, end);
    // 向上更新
    pushUp(node);
}
// 向上更新
private void pushUp(Node node) {
    node.val = node.left.val + node.right.val;
}

题目中都没有给出很具体的范围：

解法：动态开点

我们手动模拟一下过程：
栗子：nums = [1, 2, 3, 4, 5]

线段树初始状态：数组区间：[0, 4]

添加一个元素 [2, 2]; val = 3

橙色叶子节点值为0，仅仅只是被创建出来了，并无实际的值。

线段树的更新：

更新：区间更新！！！【找 + 更】

如果我们要把区间 [2, 4] 内的元素都「+1」(箭头数字代表步骤顺序)

我们使用「懒惰标记」的方法，给该区间对应的节点加一个懒惰标记，表示该节点所有对应的节点都应该有此更新（节点遍历的时候会把「懒惰标记」下推给节点）

修改一下Node数据结构：

class Node {
    // 左右***节点
    Node left, right;
    // 当前节点值
    int val;
    // 懒惰标记
    int add;
}

下推懒惰标记：

// leftNum 和 rightNum 表示左右区间的叶子节点数量
// 因为如果是「加减」更新操作的话，需要用懒惰标记的值 + 叶子节点的数量
private void pushDown(Node node, int leftNum, int rightNum) {
    // 动态开点
    if (node.left == null) node.left = new Node();
    if (node.right == null) node.right = new Node();
    // 如果 add 为 0，表示没有标记
    if (node.add == 0) return ;
    // 注意：当前节点加上标记值 + 该子树所有叶子节点的数量
    node.left.val += node.add * leftNum;
    node.right.val += node.add * rightNum;
    // 把标记下推给节点
    // 对区间进行「加减」的更新操作，下推懒惰标记时需要累加起来，不能直接覆盖
    node.left.add += node.add;
    node.right.add += node.add;
    // 取消当前节点标记
    node.add = 0;
}

更新函数：

// 在区间 [start, end] 中更新区间 [l, r] 的值，将区间 [l, r] + val
// 对于上面的例子，应该这样调用该函数：update(root, 0, 4, 2, 4, 1)
public void update(Node node, int start, int end, int l, int r, int val) {
    // 找到满足要求的区间
    if (l <= start && end <= r) {
        // 区间节点加上更新值
        // 注意：需要+该子树所有叶子节点
        node.val += (end - start + 1) * val;
        // 添加懒惰标记
        // 对区间进行「加减」的更新操作，懒惰标记需要累加，不能直接覆盖
        node.add += val;
        return ;
    }
    int mid = (start + end) >> 1;
    // 下推标记
    // mid - start + 1：表示左区间叶子节点数量
    // end - mid：表示右区间叶子节点数量
    pushDown(node, mid - start + 1, end - mid);
    // [start, mid] 和 [l, r] 可能有交集，遍历左区间
    if (l <= mid) update(node.left, start, mid, l, r, val);
    // [mid + 1, end] 和 [l, r] 可能有交集，遍历右区间
    if (r > mid) update(node.right, mid + 1, end, l, r, val);
    // 向上更新
    pushUp(node);
}

线段树的查询：
查询区间 [2, 4]

代码实现：

// 在区间 [start, end] 中查询区间 [l, r] 的结果，即 [l ,r] 保持不变
// 对于上面的例子，应该这样调用该函数：query(root, 0, 4, 2, 4)
public int query(Node node, int start, int end, int l, int r) {
    // 区间 [l ,r] 完全包含区间 [start, end]
    // 例如：[2, 4] = [2, 2] + [3, 4]，当 [start, end] = [2, 2] 或者 [start, end] = [3, 4]，直接返回
    if (l <= start && end <= r) return node.val;
    // 把当前区间 [start, end] 均分得到左右***的区间范围
    // node 左***区间 [start, mid]
    // node 右***区间 [mid + 1, end]
    int mid = (start + end) >> 1, ans = 0;
    // 下推标记
    pushDown(node, mid - start + 1, end - mid);
    // [start, mid] 和 [l, r] 可能有交集，遍历左区间
    if (l <= mid) ans += query(node.left, start, mid, l, r);
    // [mid + 1, end] 和 [l, r] 可能有交集，遍历右***区间
    if (r > mid) ans += query(node.right, mid + 1, end, l, r);
    // ans 把左右子树的结果都累加起来了，与树的后续遍历同理
    return ans;
}

线段树完整模板：

基于求「区间和」以及对区间进行「加减」的更新操作，且为「动态开点」

点更新 == 区间长度为1的区间更新

public class SegmentTreeDynamic {
    class Node {
        Node left, right;
        int val, add;
    }
    private int N = (int) 1e9;
    private Node root = new Node();
    public void update(Node node, int start, int end, int l, int r, int val) {
        if (l <= start && end <= r) {
            node.val += (end - start + 1) * val;
            node.add += val;
            return ;
        }
        int mid = (start + end) >> 1;
        pushDown(node, mid - start + 1, end - mid);
        if (l <= mid) update(node.left, start, mid, l, r, val);
        if (r > mid) update(node.right, mid + 1, end, l, r, val);
        pushUp(node);
    }
    public int query(Node node, int start, int end, int l, int r) {
        if (l <= start && end <= r) return node.val;
        int mid = (start + end) >> 1, ans = 0;
        pushDown(node, mid - start + 1, end - mid);
        if (l <= mid) ans += query(node.left, start, mid, l, r);
        if (r > mid) ans += query(node.right, mid + 1, end, l, r);
        return ans;
    }
    private void pushUp(Node node) {
        node.val = node.left.val + node.right.val;
    }
    private void pushDown(Node node, int leftNum, int rightNum) {
        if (node.left == null) node.left = new Node();
        if (node.right == null) node.right = new Node();
        if (node.add == 0) return ;
        node.left.val += node.add * leftNum;
        node.right.val += node.add * rightNum;
        // 对区间进行「加减」的更新操作，下推懒惰标记时需要累加起来，不能直接覆盖
        node.left.add += node.add;
        node.right.add += node.add;
        node.add = 0;
    }
}

三、例题

题：729. 我的日程安排表 I

实现一个 MyCalendar 类来存放你的日程安排。如果要添加的日程安排不会造成 重复预订 ，则可以存储这个新的日程安排。

当两个日程安排有一些时间上的交叉时（例如两个日程安排都在同一时间内），就会产生 重复预订 。

日程可以用一对整数 start 和 end 表示，这里的时间是半开区间，即 [start, end), 实数 x 的范围为， start <= x < end 。

实现 MyCalendar 类：

• MyCalendar() 初始化日历对象。
• boolean book(int start, int end) 如果可以将日程安排成功添加到日历中而不会导致重复预订，返回 true 。否则，返回 false 并且不要将该日程安排添加到日历中。

示例：

输入：
["MyCalendar", "book", "book", "book"]
[[], [10, 20], [15, 25], [20, 30]]
输出：
[null, true, false, true]

解释：
MyCalendar myCalendar = new MyCalendar();
myCalendar.book(10, 20); // return True
myCalendar.book(15, 25); // return False ，这个日程安排不能添加到日历中，因为时间 15 已经被另一个日程安排预订了。
myCalendar.book(20, 30); // return True ，这个日程安排可以添加到日历中，因为第一个日程安排预订的每个时间都小于 20 ，且不包含时间 20 。

提示：

0 <= start < end <= 109
每个测试用例，调用 book 方法的次数最多不超过 1000 次。

解：

AC代码：

class MyCalendar {
    // 线段树
    public MyCalendar() {}
    
    public boolean book(int start, int end) {
        if(query(root, 0, N, start, end - 1) != 0) return false;
        update(root, 0, N, start, end - 1, 1);
        return true;
    }
    class Node {
        // 左右节点，节点值，懒惰标记点
        Node left, right; int val, add;
    }
    int N = (int) 1e9;
    Node root = new Node();
    // start < l <= mid < r < end
    int query(Node node, int start, int end, int l, int r) {
        // 越界
        if(l <= start && end <= r) return node.val;
        pushDown(node);
        int mid = (start + end) >> 1, ans = 0;
        if(l <= mid) ans = query(node.left, start, mid, l, r);
        if(r > mid) ans = Math.max(ans, query(node.right, mid + 1, end, l, r));
        return ans;
    }
    void pushDown(Node node) {
        if(node.left == null) node.left =  new Node();
        if(node.right == null) node.right = new Node();
        if(node.add == 0) return;
        node.left.val += node.add; node.left.add += node.add;
        node.right.val += node.add; node.right.add += node.add;
        node.add = 0; // 取消当前节点标记
    }
    void update(Node node, int start, int end, int l, int r, int val) {
        if(l <= start && end <= r) {
            node.val += val; node.add += val; return;
        }
        pushDown(node);
        int mid = (start + end) >> 1;
        if(l <= mid) update(node.left, start, mid, l , r, val);
        if(r > mid) update(node.right, mid + 1, end, l , r, val);
        pushUp(node);
    }
    void pushUp(Node node) {
        node.val = Math.max(node.left.val, node.right.val);
    }
}