矩形和取得
(DataStructure/Wavelet/RectangleSum.py)

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

概要

二次元平面上の $n$ 個の重み付き点に対して、前計算を $O(n \log n)$ で行い、矩形範囲内の重み総和クエリに $O(\log n)$ で答えるデータ構造。内部では累積和を保持した 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 n)$

Depends on

• ビットベクトル (DataStructure/Wavelet/BitVector.py)

Verified with

• TestCase/LibraryChecker/rectangle_sum.WaveletMatrix.test.py

Code

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


class RectangleSum:
    def __init__(self, ys, vals, MAXLOG=32):
        self.MAXLOG = MAXLOG
        self.n = len(vals)
        self.mat = []
        self.mid = []
        self.data = [[0] * (self.n + 1) for i in range(self.MAXLOG)]

        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] & (1 << d):
                    rs.append(od)
                    vec.set(i)
                else:
                    ls.append(od)
            vec.build()
            self.mat.append(vec)
            self.mid.append(len(ls))
            order = ls + rs
            for i, val in enumerate(order):
                self.data[- d - 1][i + 1] = self.data[- d - 1][i] + vals[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][self.mat[d].rank(r, 0)]
                res -= self.data[d][self.mat[d].rank(l, 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 CompressedRectangleSum:
    def __init__(self, points):
        points = sorted(points)
        xs, ys, vals = zip(*points)
        self.xs = xs
        self.ys = sorted(set(ys))
        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 = RectangleSum(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)

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