セグメントツリーの基本構造と操作

セグメントツリーは、範囲クエリや更新操作を効率的に行うためのデータ構造です。以下では、いくつかの問題に対してセグメントツリーをどのように適用するかを説明します。

#include 
using namespace std;

#define FOR(i, a, b) for (int i = (a); i <= (b); ++i)
#define ROF(i, a, b) for (int i = (a); i >= (b); --i)
#define DEBUG(x) cerr << #x << " = " << x << endl
typedef long long ll;
typedef pair PII;

namespace SegmentTree {
    const int N = 100005;
    struct Node {
        int l, r;
        ll sum, lazy_add;
    } tree[N * 4];
    ll data[N];

    void push_up(int u) {
        tree[u].sum = tree[u * 2].sum + tree[u * 2 + 1].sum;
    }

    void push_down(int u) {
        if (tree[u].lazy_add != 0) {
            tree[u * 2].sum += tree[u].lazy_add * (tree[u * 2].r - tree[u * 2].l + 1);
            tree[u * 2 + 1].sum += tree[u].lazy_add * (tree[u * 2 + 1].r - tree[u * 2 + 1].l + 1);
            tree[u * 2].lazy_add += tree[u].lazy_add;
            tree[u * 2 + 1].lazy_add += tree[u].lazy_add;
            tree[u].lazy_add = 0;
        }
    }

    void build(int u, int l, int r) {
        tree[u].l = l; tree[u].r = r;
        tree[u].lazy_add = 0;
        if (l == r) {
            tree[u].sum = data[l];
            return;
        }
        int mid = (l + r) / 2;
        build(u * 2, l, mid);
        build(u * 2 + 1, mid + 1, r);
        push_up(u);
    }

    void update(int u, int ql, int qr, ll val) {
        if (ql <= tree[u].l && tree[u].r <= qr) {
            tree[u].sum += val * (tree[u].r - tree[u].l + 1);
            tree[u].lazy_add += val;
            return;
        }
        push_down(u);
        int mid = (tree[u].l + tree[u].r) / 2;
        if (ql <= mid) update(u * 2, ql, qr, val);
        if (qr > mid) update(u * 2 + 1, ql, qr, val);
        push_up(u);
    }

    ll query(int u, int ql, int qr) {
        if (ql <= tree[u].l && tree[u].r <= qr) {
            return tree[u].sum;
        }
        push_down(u);
        int mid = (tree[u].l + tree[u].r) / 2;
        ll res = 0;
        if (ql <= mid) res += query(u * 2, ql, qr);
        if (qr > mid) res += query(u * 2 + 1, ql, qr);
        return res;
    }
}

void solve() {
    int n, m;
    cin >> n >> m;
    for (int i = 1; i <= n; ++i) {
        cin >> SegmentTree::data[i];
    }
    SegmentTree::build(1, 1, n);
    while (m--) {
        int type, x, y;
        ll k;
        cin >> type >> x >> y;
        if (type == 1) {
            cin >> k;
            SegmentTree::update(1, x, y, k);
        } else {
            cout << SegmentTree::query(1, x, y) << "\n";
        }
    }
}

複数の区間操作をサポートしたセグメントツリー

次に、加算と乗算の両方をサポートするセグメントツリーについて説明します。

namespace AdvancedSegmentTree {
    const int N = 100005;
    struct Node {
        int l, r;
        ll sum, add_lazy, mul_lazy;
    } tree[N * 4];
    ll data[N], mod;

    void push_up(int u) {
        tree[u].sum = (tree[u * 2].sum + tree[u * 2 + 1].sum) % mod;
    }

    void push_down(int u) {
        if (tree[u].mul_lazy != 1) {
            tree[u * 2].sum = (tree[u * 2].sum * tree[u].mul_lazy) % mod;
            tree[u * 2 + 1].sum = (tree[u * 2 + 1].sum * tree[u].mul_lazy) % mod;
            tree[u * 2].add_lazy = (tree[u * 2].add_lazy * tree[u].mul_lazy) % mod;
            tree[u * 2 + 1].add_lazy = (tree[u * 2 + 1].add_lazy * tree[u].mul_lazy) % mod;
            tree[u * 2].mul_lazy = (tree[u * 2].mul_lazy * tree[u].mul_lazy) % mod;
            tree[u * 2 + 1].mul_lazy = (tree[u * 2 + 1].mul_lazy * tree[u].mul_lazy) % mod;
        }
        if (tree[u].add_lazy != 0) {
            tree[u * 2].sum = (tree[u * 2].sum + tree[u].add_lazy * (tree[u * 2].r - tree[u * 2].l + 1)) % mod;
            tree[u * 2 + 1].sum = (tree[u * 2 + 1].sum + tree[u].add_lazy * (tree[u * 2 + 1].r - tree[u * 2 + 1].l + 1)) % mod;
            tree[u * 2].add_lazy = (tree[u * 2].add_lazy + tree[u].add_lazy) % mod;
            tree[u * 2 + 1].add_lazy = (tree[u * 2 + 1].add_lazy + tree[u].add_lazy) % mod;
        }
        tree[u].add_lazy = 0;
        tree[u].mul_lazy = 1;
    }

    void build(int u, int l, int r) {
        tree[u].l = l; tree[u].r = r;
        tree[u].add_lazy = 0; tree[u].mul_lazy = 1;
        if (l == r) {
            tree[u].sum = data[l];
            return;
        }
        int mid = (l + r) / 2;
        build(u * 2, l, mid);
        build(u * 2 + 1, mid + 1, r);
        push_up(u);
    }

    void update_add(int u, int ql, int qr, ll val) {
        if (ql <= tree[u].l && tree[u].r <= qr) {
            tree[u].sum = (tree[u].sum + val * (tree[u].r - tree[u].l + 1)) % mod;
            tree[u].add_lazy = (tree[u].add_lazy + val) % mod;
            return;
        }
        push_down(u);
        int mid = (tree[u].l + tree[u].r) / 2;
        if (ql <= mid) update_add(u * 2, ql, qr, val);
        if (qr > mid) update_add(u * 2 + 1, ql, qr, val);
        push_up(u);
    }

    void update_mul(int u, int ql, int qr, ll val) {
        if (ql <= tree[u].l && tree[u].r <= qr) {
            tree[u].sum = (tree[u].sum * val) % mod;
            tree[u].add_lazy = (tree[u].add_lazy * val) % mod;
            tree[u].mul_lazy = (tree[u].mul_lazy * val) % mod;
            return;
        }
        push_down(u);
        int mid = (tree[u].l + tree[u].r) / 2;
        if (ql <= mid) update_mul(u * 2, ql, qr, val);
        if (qr > mid) update_mul(u * 2 + 1, ql, qr, val);
        push_up(u);
    }

    ll query(int u, int ql, int qr) {
        if (ql <= tree[u].l && tree[u].r <= qr) {
            return tree[u].sum;
        }
        push_down(u);
        int mid = (tree[u].l + tree[u].r) / 2;
        ll res = 0;
        if (ql <= mid) res = (res + query(u * 2, ql, qr)) % mod;
        if (qr > mid) res = (res + query(u * 2 + 1, ql, qr)) % mod;
        return res;
    }
}

void advanced_solve() {
    int n, q;
    cin >> n >> q >> AdvancedSegmentTree::mod;
    for (int i = 1; i <= n; ++i) {
        cin >> AdvancedSegmentTree::data[i];
    }
    AdvancedSegmentTree::build(1, 1, n);
    while (q--) {
        int type, x, y;
        ll k;
        cin >> type >> x >> y;
        if (type == 1) {
            cin >> k;
            AdvancedSegmentTree::update_mul(1, x, y, k);
        } else if (type == 2) {
            cin >> k;
            AdvancedSegmentTree::update_add(1, x, y, k);
        } else {
            cout << AdvancedSegmentTree::query(1, x, y) << "\n";
        }
    }
}

他の応用例

ここでは、区間の最大値や最小値を求めるためにセグメントツリーを使用する方法についても紹介します。

namespace MaxSegmentTree {
    const int N = 100005;
    struct Node {
        int l, r;
        ll max_val;
    } tree[N * 4];
    ll data[N];

    void push_up(int u) {
        tree[u].max_val = max(tree[u * 2].max_val, tree[u * 2 + 1].max_val);
    }

    void build(int u, int l, int r) {
        tree[u].l = l; tree[u].r = r;
        if (l == r) {
            tree[u].max_val = data[l];
            return;
        }
        int mid = (l + r) / 2;
        build(u * 2, l, mid);
        build(u * 2 + 1, mid + 1, r);
        push_up(u);
    }

    void update(int u, int pos, ll val) {
        if (tree[u].l == tree[u].r) {
            tree[u].max_val = val;
            return;
        }
        int mid = (tree[u].l + tree[u].r) / 2;
        if (pos <= mid) update(u * 2, pos, val);
        else update(u * 2 + 1, pos, val);
        push_up(u);
    }

    ll query_max(int u, int ql, int qr) {
        if (ql <= tree[u].l && tree[u].r <= qr) {
            return tree[u].max_val;
        }
        int mid = (tree[u].l + tree[u].r) / 2;
        ll res = LLONG_MIN;
        if (ql <= mid) res = max(res, query_max(u * 2, ql, qr));
        if (qr > mid) res = max(res, query_max(u * 2 + 1, ql, qr));
        return res;
    }
}