segment_tree.cpp

/*
Segment Tree
Mathematical Foundation: Binary tree for range queries/updates
Each node stores info about segment [l,r]
Range queries: O(log n), Point/Range updates: O(log n)
Space: O(4n), Build: O(n)
*/

#include <bits/stdc++.h>
using namespace std;

// Basic Segment Tree (Range Sum Query)
// Related LeetCode Problems:
// 307. Range Sum Query - Mutable
// https://leetcode.com/problems/range-sum-query-mutable/
// 315. Count of Smaller Numbers After Self
// https://leetcode.com/problems/count-of-smaller-numbers-after-self/
// 327. Count of Range Sum
// https://leetcode.com/problems/count-of-range-sum/
// 493. Reverse Pairs
// https://leetcode.com/problems/reverse-pairs/
// 673. Number of Longest Increasing Subsequence
// https://leetcode.com/problems/number-of-longest-increasing-subsequence/
// 715. Range Module
// https://leetcode.com/problems/range-module/
// 729. My Calendar I
// https://leetcode.com/problems/my-calendar-i/
// 731. My Calendar II
// https://leetcode.com/problems/my-calendar-ii/
// 732. My Calendar III
// https://leetcode.com/problems/my-calendar-iii/
// 850. Rectangle Area II
// https://leetcode.com/problems/rectangle-area-ii/
class SegmentTree {
    vector<long long> tree;
    int n;
    
public:
    SegmentTree(vector<int>& arr) {
        n = arr.size();
        tree.resize(4 * n);
        build(arr, 1, 0, n - 1);
    }
    
    void build(vector<int>& arr, int node, int start, int end) {
        if (start == end) {
            tree[node] = arr[start];
        } else {
            int mid = (start + end) / 2;
            build(arr, 2 * node, start, mid);
            build(arr, 2 * node + 1, mid + 1, end);
            tree[node] = tree[2 * node] + tree[2 * node + 1];
        }
    }
    
    void update(int node, int start, int end, int idx, int val) {
        if (start == end) {
            tree[node] = val;
        } else {
            int mid = (start + end) / 2;
            if (idx <= mid) update(2 * node, start, mid, idx, val);
            else update(2 * node + 1, mid + 1, end, idx, val);
            tree[node] = tree[2 * node] + tree[2 * node + 1];
        }
    }
    
    long long query(int node, int start, int end, int l, int r) {
        if (r < start || end < l) return 0;
        if (l <= start && end <= r) return tree[node];
        int mid = (start + end) / 2;
        return query(2 * node, start, mid, l, r) + 
               query(2 * node + 1, mid + 1, end, l, r);
    }
    
    void update(int idx, int val) { update(1, 0, n - 1, idx, val); }
    long long query(int l, int r) { return query(1, 0, n - 1, l, r); }
};

// Lazy Propagation Segment Tree
// Related LeetCode Problems:
// 699. Falling Squares
// https://leetcode.com/problems/falling-squares/
// 715. Range Module
// https://leetcode.com/problems/range-module/
// 729. My Calendar I
// https://leetcode.com/problems/my-calendar-i/
// 731. My Calendar II
// https://leetcode.com/problems/my-calendar-ii/
// 732. My Calendar III
// https://leetcode.com/problems/my-calendar-iii/
// 850. Rectangle Area II
// https://leetcode.com/problems/rectangle-area-ii/
// 1157. Online Majority Element In Subarray
// https://leetcode.com/problems/online-majority-element-in-subarray/
// 1505. Minimum Possible Integer After at Most K Adjacent Swaps On Digits
// https://leetcode.com/problems/minimum-possible-integer-after-at-most-k-adjacent-swaps-on-digits/
// 1622. Fancy Sequence
// https://leetcode.com/problems/fancy-sequence/
// 2286. Booking Concert Tickets in Groups
// https://leetcode.com/problems/booking-concert-tickets-in-groups/
class LazySegmentTree {
    vector<long long> tree, lazy;
    int n;
    
public:
    LazySegmentTree(vector<int>& arr) {
        n = arr.size();
        tree.resize(4 * n);
        lazy.resize(4 * n);
        build(arr, 1, 0, n - 1);
    }
    
    void build(vector<int>& arr, int node, int start, int end) {
        if (start == end) {
            tree[node] = arr[start];
        } else {
            int mid = (start + end) / 2;
            build(arr, 2 * node, start, mid);
            build(arr, 2 * node + 1, mid + 1, end);
            tree[node] = tree[2 * node] + tree[2 * node + 1];
        }
    }
    
    void push(int node, int start, int end) {
        if (lazy[node] != 0) {
            tree[node] += lazy[node] * (end - start + 1);
            if (start != end) {
                lazy[2 * node] += lazy[node];
                lazy[2 * node + 1] += lazy[node];
            }
            lazy[node] = 0;
        }
    }
    
    void updateRange(int node, int start, int end, int l, int r, int val) {
        push(node, start, end);
        if (start > r || end < l) return;
        if (start >= l && end <= r) {
            lazy[node] += val;
            push(node, start, end);
            return;
        }
        int mid = (start + end) / 2;
        updateRange(2 * node, start, mid, l, r, val);
        updateRange(2 * node + 1, mid + 1, end, l, r, val);
        push(2 * node, start, mid);
        push(2 * node + 1, mid + 1, end);
        tree[node] = tree[2 * node] + tree[2 * node + 1];
    }
    
    long long queryRange(int node, int start, int end, int l, int r) {
        if (start > r || end < l) return 0;
        push(node, start, end);
        if (start >= l && end <= r) return tree[node];
        int mid = (start + end) / 2;
        return queryRange(2 * node, start, mid, l, r) + 
               queryRange(2 * node + 1, mid + 1, end, l, r);
    }
    
    void updateRange(int l, int r, int val) { updateRange(1, 0, n - 1, l, r, val); }
    long long queryRange(int l, int r) { return queryRange(1, 0, n - 1, l, r); }
};

// Min/Max Segment Tree
// Related LeetCode Problems:
// 239. Sliding Window Maximum
// https://leetcode.com/problems/sliding-window-maximum/
// 224. Basic Calculator
// https://leetcode.com/problems/basic-calculator/
// 315. Count of Smaller Numbers After Self
// https://leetcode.com/problems/count-of-smaller-numbers-after-self/
// 493. Reverse Pairs
// https://leetcode.com/problems/reverse-pairs/
// 699. Falling Squares
// https://leetcode.com/problems/falling-squares/
// 715. Range Module
// https://leetcode.com/problems/range-module/
// 850. Rectangle Area II
// https://leetcode.com/problems/rectangle-area-ii/
// 1157. Online Majority Element In Subarray
// https://leetcode.com/problems/online-majority-element-in-subarray/
// 1505. Minimum Possible Integer After at Most K Adjacent Swaps On Digits
// https://leetcode.com/problems/minimum-possible-integer-after-at-most-k-adjacent-swaps-on-digits/
// 2286. Booking Concert Tickets in Groups
// https://leetcode.com/problems/booking-concert-tickets-in-groups/
class MinMaxSegTree {
    vector<int> minTree, maxTree;
    int n;
    
public:
    MinMaxSegTree(vector<int>& arr) {
        n = arr.size();
        minTree.resize(4 * n);
        maxTree.resize(4 * n);
        build(arr, 1, 0, n - 1);
    }
    
    void build(vector<int>& arr, int node, int start, int end) {
        if (start == end) {
            minTree[node] = maxTree[node] = arr[start];
        } else {
            int mid = (start + end) / 2;
            build(arr, 2 * node, start, mid);
            build(arr, 2 * node + 1, mid + 1, end);
            minTree[node] = min(minTree[2 * node], minTree[2 * node + 1]);
            maxTree[node] = max(maxTree[2 * node], maxTree[2 * node + 1]);
        }
    }
    
    void update(int node, int start, int end, int idx, int val) {
        if (start == end) {
            minTree[node] = maxTree[node] = val;
        } else {
            int mid = (start + end) / 2;
            if (idx <= mid) update(2 * node, start, mid, idx, val);
            else update(2 * node + 1, mid + 1, end, idx, val);
            minTree[node] = min(minTree[2 * node], minTree[2 * node + 1]);
            maxTree[node] = max(maxTree[2 * node], maxTree[2 * node + 1]);
        }
    }
    
    int queryMin(int node, int start, int end, int l, int r) {
        if (r < start || end < l) return INT_MAX;
        if (l <= start && end <= r) return minTree[node];
        int mid = (start + end) / 2;
        return min(queryMin(2 * node, start, mid, l, r),
                  queryMin(2 * node + 1, mid + 1, end, l, r));
    }
    
    int queryMax(int node, int start, int end, int l, int r) {
        if (r < start || end < l) return INT_MIN;
        if (l <= start && end <= r) return maxTree[node];
        int mid = (start + end) / 2;
        return max(queryMax(2 * node, start, mid, l, r),
                  queryMax(2 * node + 1, mid + 1, end, l, r));
    }
    
    void update(int idx, int val) { update(1, 0, n - 1, idx, val); }
    int queryMin(int l, int r) { return queryMin(1, 0, n - 1, l, r); }
    int queryMax(int l, int r) { return queryMax(1, 0, n - 1, l, r); }
};

// Count of elements <= x in range [l,r]
// Related LeetCode Problems:
// 315. Count of Smaller Numbers After Self
// https://leetcode.com/problems/count-of-smaller-numbers-after-self/
// 327. Count of Range Sum
// https://leetcode.com/problems/count-of-range-sum/
// 493. Reverse Pairs
// https://leetcode.com/problems/reverse-pairs/
// 673. Number of Longest Increasing Subsequence
// https://leetcode.com/problems/number-of-longest-increasing-subsequence/
// 1157. Online Majority Element In Subarray
// https://leetcode.com/problems/online-majority-element-in-subarray/
// 1395. Count Number of Teams
// https://leetcode.com/problems/count-number-of-teams/
// 1505. Minimum Possible Integer After at Most K Adjacent Swaps On Digits
// https://leetcode.com/problems/minimum-possible-integer-after-at-most-k-adjacent-swaps-on-digits/
// 1649. Create Sorted Array through Instructions
// https://leetcode.com/problems/create-sorted-array-through-instructions/
// 2179. Count Good Triplets in an Array
// https://leetcode.com/problems/count-good-triplets-in-an-array/
class CountSegTree {
    vector<vector<int>> tree;
    int n;
    
public:
    CountSegTree(vector<int>& arr) {
        n = arr.size();
        tree.resize(4 * n);
        build(arr, 1, 0, n - 1);
    }
    
    void build(vector<int>& arr, int node, int start, int end) {
        if (start == end) {
            tree[node] = {arr[start]};
        } else {
            int mid = (start + end) / 2;
            build(arr, 2 * node, start, mid);
            build(arr, 2 * node + 1, mid + 1, end);
            merge(tree[2 * node].begin(), tree[2 * node].end(),
                  tree[2 * node + 1].begin(), tree[2 * node + 1].end(),
                  back_inserter(tree[node]));
        }
    }
    
    int query(int node, int start, int end, int l, int r, int x) {
        if (r < start || end < l) return 0;
        if (l <= start && end <= r) {
            return upper_bound(tree[node].begin(), tree[node].end(), x) - tree[node].begin();
        }
        int mid = (start + end) / 2;
        return query(2 * node, start, mid, l, r, x) + 
               query(2 * node + 1, mid + 1, end, l, r, x);
    }
    
    int query(int l, int r, int x) { return query(1, 0, n - 1, l, r, x); }
};