暫存一下線段樹模板

potential-star發表於2024-10-30
template <class Info>  // 模板類,Info 是一個模板引數,表示線段樹節點儲存的資訊型別
struct SegmentTree {
    int n;              // 表示線段樹中儲存的元素個數
    vector<Info> info;  // 用於儲存線段樹節點的陣列,型別為 Info

    SegmentTree() : n(0) {}  // 預設建構函式,初始化 n 為 0
    // 傳入 n_ 表示元素個數,v_ 表示初始化的預設值
    SegmentTree(int n_, Info v_ = Info()) {
        init(n_, v_);
    }

    // 模板建構函式:傳入一個向量,用於初始化線段樹
    template <class T>
    SegmentTree(vector<T> init_) {
        init(init_);
    }

    // 初始化函式:傳入元素個數和預設值
    void init(int n_, Info v_ = Info()) {
        init(vector(n_, v_));  // 用一個大小為 n_ 的向量初始化,元素值為 v_
    }

    // 模板初始化函式:傳入一個向量,用於初始化線段樹
    template <class T>
    void init(vector<T> init_) {
        n = init_.size();  // 記錄元素個數
        // 初始化 info 陣列的大小,4 << __lg(n) 是 2 的 log(n) 次方乘以 4
        info.assign(4 << __lg(n), Info());
        auto build = [&](auto self, int p, int l, int r) -> void {
            if (r - l == 1) {        // 如果區間長度為 1,說明是葉節點
                info[p] = init_[l];  // 將葉節點賦值為初始化向量的對應值
                return;
            }
            int m = (l + r) / 2;          // 計算區間中點
            self(self, 2 * p, l, m);      // 遞迴構建左子樹
            self(self, 2 * p + 1, m, r);  // 遞迴構建右子樹
            pushup(p);                    // 更新當前節點的值為其兩個子節點的合併值
        };

        build(build, 1, 0, n);  // 呼叫 Lambda 函式,從根節點開始構建線段樹
    }

    void pushup(int p) {
        info[p] = info[2 * p] + info[2 * p + 1];
    }
    void modify(int p, int l, int r, int x, const Info &v) {
        if (r - l == 1) {  // 如果區間長度為 1,說明是葉節點
            info[p] = v;   // 直接修改
            return;
        }
        int m = (l + r) / 2;
        if (x < m) {
            modify(2 * p, l, m, x, v);  // 遞迴修改左子樹
        } else {
            modify(2 * p + 1, m, r, x, v);  // 遞迴修改右子樹
        }
        pushup(p);  // 更新當前節點的值為其兩個子節點的合併值
    }

    // 對外的修改介面,只需傳入修改位置和新值
    void modify(int p, const Info &v) {
        modify(1, 0, n, p, v);  // 從根節點開始遞迴修改
    }

    // 區間查詢函式:遞迴查詢區間 [x, y) 的資訊
    Info rangeQuery(int p, int l, int r, int x, int y) {
        if (l >= y || r <= x) {  // 如果查詢區間與當前區間無交集
            return Info();       // 返回預設的 Info 物件
        }
        if (l >= x && r <= y) {  // 如果當前區間完全包含在查詢區間內
            return info[p];      // 返回當前節點的值
        }
        int m = (l + r) / 2;  // 計算區間中點
        // 遞迴查詢左子樹和右子樹,併合並結果
        return rangeQuery(2 * p, l, m, x, y) + rangeQuery(2 * p + 1, m, r, x, y);
    }

    // 對外的區間查詢介面,傳入區間 [l, r)
    Info rangeQuery(int l, int r) {
        return rangeQuery(1, 0, n, l, r);  // 從根節點開始查詢
    }

    // 查詢區間內第一個滿足條件的元素
    template <class F>
    int findFirst(int p, int l, int r, int x, int y, F &&pred) {
        if (l >= y || r <= x) {  // 如果查詢區間與當前區間無交集
            return -1;           // 返回 -1 表示未找到
        }
        if (l >= x && r <= y && !pred(info[p])) {  // 如果當前區間完全包含在查詢區間內且不滿足條件
            return -1;                             // 返回 -1 表示未找到
        }
        if (r - l == 1) {  // 如果區間長度為 1,說明是葉節點
            return l;      // 返回葉節點的位置
        }
        int m = (l + r) / 2;                           // 計算區間中點
        int res = findFirst(2 * p, l, m, x, y, pred);  // 遞迴查詢左子樹
        if (res == -1) {                               // 如果左子樹未找到,繼續查詢右子樹
            res = findFirst(2 * p + 1, m, r, x, y, pred);
        }
        return res;  // 返回結果
    }

    // 對外的查詢介面,查詢區間 [l, r) 內第一個滿足條件的元素
    template <class F>
    int findFirst(int l, int r, F &&pred) {
        return findFirst(1, 0, n, l, r, pred);  // 從根節點開始查詢
    }

    // 查詢區間內最後一個滿足條件的元素
    template <class F>
    int findLast(int p, int l, int r, int x, int y, F &&pred) {
        if (l >= y || r <= x) {  // 如果查詢區間與當前區間無交集
            return -1;           // 返回 -1 表示未找到
        }
        if (l >= x && r <= y && !pred(info[p])) {  // 如果當前區間完全包含在查詢區間內且不滿足條件
            return -1;                             // 返回 -1 表示未找到
        }
        if (r - l == 1) {  // 如果區間長度為 1,說明是葉節點
            return l;      // 返回葉節點的位置
        }
        int m = (l + r) / 2;
        int res = findLast(2 * p + 1, m, r, x, y, pred);  // 遞迴查詢右子樹
        if (res == -1) {                                  // 如果右子樹未找到,繼續查詢左子樹
            res = findLast(2 * p, l, m, x, y, pred);
        }
        return res;  // 返回結果
    }

    // 對外的查詢介面,查詢區間 [l, r) 內最後一個滿足條件的元素
    template <class F>
    int findLast(int l, int r, F &&pred) {
        return findLast(1, 0, n, l, r, pred);  // 從根節點開始查詢
    }
};