一点加算・矩形和取得
(DataStructure/Wavelet/PointAddRectangleSum.py)

• View this file on GitHub
• Last update: 2022-08-07 15:04:38+09:00

概要

二次元平面上の $n$ 個の重み付き点に対して、前計算を $O(n \log n)$ で行い、一点重み加算クエリに $O(\log^2n)$、矩形範囲内の重み総和クエリに $O(\log^2n)$ で答えるデータ構造。内部では Binary Index Tree を保持した Wavelet Matrix を構築する。

使い方

CompressedRectangleSum(points: Iterable[Tuple[int, int, int]])
二次元平面上の $n$ 個の重み付き点 points から Wavelet Matrix を構築する。各点は ($x$ 座標, $y$ 座標, 重み) で与えられる。計算量 $O(n \log n)$

• rect_sum(l: int, r: int, lower: int, upper: int) -> int
  矩形範囲 $\lbrack l, r) × \lbrack lower, upper)$ に含まれる点の重みの総和を返す。計算量 $O(\log^2n)$

• point_add(x: int, y: int, val: int) -> int
  座標 $(x, y)$ 上の点に対して、重みを val 追加する。このときの点は、予め構築の際に points の $1$ つとして与えておく必要がある。計算量 $O(\log^2n)$

Depends on

• 一点加算・区間和取得 (DataStructure/BinaryIndexedTree/PointAddRangeSum.py)
• ビットベクトル (DataStructure/Wavelet/BitVector.py)

Verified with

• TestCase/LibraryChecker/point_add_rectangle_sum.test.py

Code

from DataStructure.Wavelet.BitVector import BitVector
from DataStructure.BinaryIndexedTree.PointAddRangeSum import BinaryIndexedTree
from bisect import bisect_left


class PointAddRectangleSum:
    def __init__(self, ys, vals, MAXLOG=32):
        self.MAXLOG = MAXLOG
        self.n = len(ys)
        self.ys = ys
        self.mat = []
        self.mid = []
        self.data = []

        order = [i for i in range(self.n)]
        for d in reversed(range(self.MAXLOG)):
            vec = BitVector(self.n + 1)
            ls = []
            rs = []
            for i, od in enumerate(order):
                if (ys[od] >> d) & 1:
                    rs.append(od)
                    vec.set(i)
                else:
                    ls.append(od)
            vec.build()
            self.mat.append(vec)
            self.mid.append(len(ls))
            order = ls + rs
            self.data.append(BinaryIndexedTree(self.n))
            self.data[-1].build([vals[od] for od in order])

    def point_add(self, k, val):
        y = self.ys[k]
        for d in range(self.MAXLOG):
            if y >> (self.MAXLOG - d - 1) & 1:
                k = self.mat[d].rank(k, 1) + self.mid[d]
            else:
                k = self.mat[d].rank(k, 0)
            self.data[d].add(k, val)

    def rect_sum(self, l, r, upper):
        res = 0
        for d in range(self.MAXLOG):
            if upper >> (self.MAXLOG - d - 1) & 1:
                res += self.data[d].sum(self.mat[d].rank(l, 0), self.mat[d].rank(r, 0))
                l = self.mat[d].rank(l, 1) + self.mid[d]
                r = self.mat[d].rank(r, 1) + self.mid[d]
            else:
                l = self.mat[d].rank(l, 0)
                r = self.mat[d].rank(r, 0)
        return res


class CompressedPointAddRectangleSum:
    def __init__(self, points):
        points = sorted(points)
        xs, ys, vals = zip(*points)
        self.xs = xs
        self.ys = sorted(set(ys))
        self.idxs = {(x, y): idx for idx, (x, y, _) in enumerate(points)}
        self.comp = {val: idx for idx, val in enumerate(self.ys)}
        ys = [self.comp[val] for val in ys]
        MAXLOG = len(self.ys).bit_length()
        self.mat = PointAddRectangleSum(ys, vals, MAXLOG)

    def rect_sum(self, l, r, lower, upper):
        l = bisect_left(self.xs, l)
        r = bisect_left(self.xs, r)
        lower = bisect_left(self.ys, lower)
        upper = bisect_left(self.ys, upper)
        return self.mat.rect_sum(l, r, upper) - self.mat.rect_sum(l, r, lower)

    def point_add(self, x, y, val):
        if (x, y) not in self.idxs:
            raise KeyError(f'point(x={x}, y={y}) must be pre-given as an argument')
        idx = self.idxs[x, y]
        self.mat.point_add(idx, val)

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.12.4/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/python.py", line 96, in bundle
    raise NotImplementedError
NotImplementedError

Back to top page