2D Segment Tree
(data-structure/segtree/segment_tree_2d.hpp)

• View this file on GitHub
• Last update: 2024-03-02 18:46:36+09:00
• Include: #include "data-structure/segtree/segment_tree_2d.hpp"

Description

2D セグメント木は，モノイド $(T, \cdot, e)$ の重みを持つ2次元平面上の点集合に対する一点更新と矩形領域積取得を提供するデータ構造である．

この実装では，重みをもたせる点を先読みして初期化時に渡す必要がある．

空間計算量: $O(n)$

Operations

• SegmentTree2D(vector<pair<X, Y>> pts)
  • pts の点に対する 2D セグメント木を初期化する
  • 時間計算量: $O(n\log n)$
• void update(X x, Y y, T val)
  • 点 $(x, y)$ の重みを $val$ に更新する
  • 時間計算量: $O((\log n)^2)$
• T fold(X sx, X tx, Y sy, Y ty)
  • 矩形領域 $[s_x, t_x) \times [s_y, t_y)$ 内の点の重みの積を取得する
  • 時間計算量: $O((\log n)^2)$

Depends on

• Segment Tree (data-structure/segtree/segment_tree.hpp)

Verified with

• test/yosupo/point_add_rectangle_sum.2d_segtree.test.cpp

Code

#pragma once
#include <algorithm>
#include <bit>
#include <cassert>
#include <utility>
#include <vector>

#include "segment_tree.hpp"

template <typename X, typename Y, typename M>
class SegmentTree2D {
    using T = M::T;

   public:
    SegmentTree2D() = default;
    explicit SegmentTree2D(const std::vector<std::pair<X, Y>>& pts) {
        for (auto& [x, y] : pts) {
            xs.push_back(x);
        }
        std::ranges::sort(xs);
        xs.erase(std::ranges::unique(xs).begin(), xs.end());

        const int n = xs.size();
        size = std::bit_ceil((unsigned int)n);
        ys.resize(2 * size);
        seg.resize(2 * size);

        for (auto& [x, y] : pts) {
            ys[size + getx(x)].push_back(y);
        }

        for (int i = size + n - 1; i > 0; --i) {
            if (i >= size) {
                std::ranges::sort(ys[i]);
            } else {
                std::ranges::merge(ys[2 * i], ys[2 * i + 1],
                                   std::back_inserter(ys[i]));
            }
            ys[i].erase(std::ranges::unique(ys[i]).begin(), ys[i].end());
        }
        for (int i = 1; i < size + n; ++i) {
            seg[i] = SegmentTree<M>(ys[i].size());
        }
    }

    T get(X x, Y y) const {
        int kx = getx(x);
        assert(kx < (int)xs.size() && xs[kx] == x);
        kx += size;
        int ky = gety(kx, y);
        assert(ky < (int)ys[kx].size() && ys[kx][ky] == y);
        return seg[kx][ky];
    }

    void update(X x, Y y, T val) {
        int kx = getx(x);
        assert(kx < (int)xs.size() && xs[kx] == x);
        kx += size;
        int ky = gety(kx, y);
        assert(ky < (int)ys[kx].size() && ys[kx][ky] == y);
        seg[kx].update(ky, val);
        while (kx >>= 1) {
            ky = gety(kx, y);
            int kl = gety(2 * kx, y), kr = gety(2 * kx + 1, y);
            T vl = (kl < (int)ys[2 * kx].size() && ys[2 * kx][kl] == y
                        ? seg[2 * kx][kl]
                        : M::id());
            T vr = (kr < (int)ys[2 * kx + 1].size() && ys[2 * kx + 1][kr] == y
                        ? seg[2 * kx + 1][kr]
                        : M::id());
            seg[kx].update(ky, M::op(vl, vr));
        }
    }

    T fold(X sx, X tx, Y sy, Y ty) const {
        T ret = M::id();
        for (int l = size + getx(sx), r = size + getx(tx); l < r;
             l >>= 1, r >>= 1) {
            if (l & 1) {
                ret = M::op(ret, seg[l].fold(gety(l, sy), gety(l, ty)));
                ++l;
            }
            if (r & 1) {
                --r;
                ret = M::op(ret, seg[r].fold(gety(r, sy), gety(r, ty)));
            }
        }
        return ret;
    }

   private:
    int size;
    std::vector<X> xs;
    std::vector<std::vector<Y>> ys;
    std::vector<SegmentTree<M>> seg;

    int getx(X x) const { return std::ranges::lower_bound(xs, x) - xs.begin(); }
    int gety(int k, Y y) const {
        return std::ranges::lower_bound(ys[k], y) - ys[k].begin();
    }
};

#line 2 "data-structure/segtree/segment_tree_2d.hpp"
#include <algorithm>
#include <bit>
#include <cassert>
#include <utility>
#include <vector>

#line 5 "data-structure/segtree/segment_tree.hpp"

template <typename M>
class SegmentTree {
    using T = M::T;

   public:
    SegmentTree() = default;
    explicit SegmentTree(int n) : SegmentTree(std::vector<T>(n, M::id())) {}
    explicit SegmentTree(const std::vector<T>& v)
        : size(std::bit_ceil(v.size())), node(2 * size, M::id()) {
        std::ranges::copy(v, node.begin() + size);
        for (int i = size - 1; i > 0; --i) {
            node[i] = M::op(node[2 * i], node[2 * i + 1]);
        }
    }

    T operator[](int k) const { return node[k + size]; }

    void update(int k, const T& x) {
        k += size;
        node[k] = x;
        while (k >>= 1) node[k] = M::op(node[2 * k], node[2 * k + 1]);
    }

    T fold(int l, int r) const {
        T vl = M::id(), vr = M::id();
        for (l += size, r += size; l < r; l >>= 1, r >>= 1) {
            if (l & 1) vl = M::op(vl, node[l++]);
            if (r & 1) vr = M::op(node[--r], vr);
        }
        return M::op(vl, vr);
    }

    template <typename F>
    int find_first(int l, F cond) const {
        T v = M::id();
        for (l += size; l > 0; l >>= 1) {
            if (l & 1) {
                T nv = M::op(v, node[l]);
                if (cond(nv)) {
                    while (l < size) {
                        nv = M::op(v, node[2 * l]);
                        if (cond(nv)) {
                            l = 2 * l;
                        } else {
                            v = nv, l = 2 * l + 1;
                        }
                    }
                    return l + 1 - size;
                }
                v = nv;
                ++l;
            }
        }
        return -1;
    }

    template <typename F>
    int find_last(int r, F cond) const {
        T v = M::id();
        for (r += size; r > 0; r >>= 1) {
            if (r & 1) {
                --r;
                T nv = M::op(node[r], v);
                if (cond(nv)) {
                    while (r < size) {
                        nv = M::op(node[2 * r + 1], v);
                        if (cond(nv)) {
                            r = 2 * r + 1;
                        } else {
                            v = nv, r = 2 * r;
                        }
                    }
                    return r - size;
                }
                v = nv;
            }
        }
        return -1;
    }

   private:
    int size;
    std::vector<T> node;
};
#line 9 "data-structure/segtree/segment_tree_2d.hpp"

template <typename X, typename Y, typename M>
class SegmentTree2D {
    using T = M::T;

   public:
    SegmentTree2D() = default;
    explicit SegmentTree2D(const std::vector<std::pair<X, Y>>& pts) {
        for (auto& [x, y] : pts) {
            xs.push_back(x);
        }
        std::ranges::sort(xs);
        xs.erase(std::ranges::unique(xs).begin(), xs.end());

        const int n = xs.size();
        size = std::bit_ceil((unsigned int)n);
        ys.resize(2 * size);
        seg.resize(2 * size);

        for (auto& [x, y] : pts) {
            ys[size + getx(x)].push_back(y);
        }

        for (int i = size + n - 1; i > 0; --i) {
            if (i >= size) {
                std::ranges::sort(ys[i]);
            } else {
                std::ranges::merge(ys[2 * i], ys[2 * i + 1],
                                   std::back_inserter(ys[i]));
            }
            ys[i].erase(std::ranges::unique(ys[i]).begin(), ys[i].end());
        }
        for (int i = 1; i < size + n; ++i) {
            seg[i] = SegmentTree<M>(ys[i].size());
        }
    }

    T get(X x, Y y) const {
        int kx = getx(x);
        assert(kx < (int)xs.size() && xs[kx] == x);
        kx += size;
        int ky = gety(kx, y);
        assert(ky < (int)ys[kx].size() && ys[kx][ky] == y);
        return seg[kx][ky];
    }

    void update(X x, Y y, T val) {
        int kx = getx(x);
        assert(kx < (int)xs.size() && xs[kx] == x);
        kx += size;
        int ky = gety(kx, y);
        assert(ky < (int)ys[kx].size() && ys[kx][ky] == y);
        seg[kx].update(ky, val);
        while (kx >>= 1) {
            ky = gety(kx, y);
            int kl = gety(2 * kx, y), kr = gety(2 * kx + 1, y);
            T vl = (kl < (int)ys[2 * kx].size() && ys[2 * kx][kl] == y
                        ? seg[2 * kx][kl]
                        : M::id());
            T vr = (kr < (int)ys[2 * kx + 1].size() && ys[2 * kx + 1][kr] == y
                        ? seg[2 * kx + 1][kr]
                        : M::id());
            seg[kx].update(ky, M::op(vl, vr));
        }
    }

    T fold(X sx, X tx, Y sy, Y ty) const {
        T ret = M::id();
        for (int l = size + getx(sx), r = size + getx(tx); l < r;
             l >>= 1, r >>= 1) {
            if (l & 1) {
                ret = M::op(ret, seg[l].fold(gety(l, sy), gety(l, ty)));
                ++l;
            }
            if (r & 1) {
                --r;
                ret = M::op(ret, seg[r].fold(gety(r, sy), gety(r, ty)));
            }
        }
        return ret;
    }

   private:
    int size;
    std::vector<X> xs;
    std::vector<std::vector<Y>> ys;
    std::vector<SegmentTree<M>> seg;

    int getx(X x) const { return std::ranges::lower_bound(xs, x) - xs.begin(); }
    int gety(int k, Y y) const {
        return std::ranges::lower_bound(ys[k], y) - ys[k].begin();
    }
};

Back to top page