二次元領域探索 (kD-Tree)
(DataStructure/misc/KDTree2D.py)

• View this file on GitHub
• Last update: 2021-02-07 18:49:02+09:00

概要

kD-Tree: 二次元平面上の点集合から領域探索を行うデータ構造

使い方

KDTree2D(points: Iterable[Tuple[int, int]])
二次元平面上の n 個の点の集合 points に対し、kD-Tree を構築する。計算量 $O(n \log^2 n)$

• find(xl: int, xr: int, yl: int, yr: int) -> List[Tuple[int, int]]
  矩形範囲 $\lbrack xl, xr) × \lbrack yl, yr)$ に含まれている点を列挙する。矩形内に含まれる点の数を k とすると、計算量 $O(\sqrt n + k)$

Verified with

• TestCase/AOJ/DSL_2_C.test.py

Code

class KDTree2D:
    def __init__(self, points):
        self.points = [tuple(p) for p in points]
        self.n = len(points)
        self.l = [-1] * self.n
        self.r = [-1] * self.n
        self.root = self._build(0, self.n, 0)

    def _build(self, l, r, depth):
        if l >= r:
            return -1
        mid = (l + r) // 2
        if depth % 2 == 0:
            self.points[l:r] = sorted(self.points[l:r], key=lambda x: x[0])
        else:
            self.points[l:r] = sorted(self.points[l:r], key=lambda x: x[1])
        self.l[mid] = self._build(l, mid, depth + 1)
        self.r[mid] = self._build(mid + 1, r, depth + 1)
        return mid

    def _find(self, xl, xr, yl, yr, depth, idx, ans):
        if idx is None:
            idx = self.root
        x, y = self.points[idx]
        if xl <= x < xr and yl <= y < yr:
            ans.append(self.points[idx])
        if depth % 2 == 0:
            if self.l[idx] != -1 and xl <= x:
                self._find(xl, xr, yl, yr, depth + 1, self.l[idx], ans)
            if self.r[idx] != -1 and x < xr:
                self._find(xl, xr, yl, yr, depth + 1, self.r[idx], ans)
        else:
            if self.l[idx] != -1 and yl <= y:
                self._find(xl, xr, yl, yr, depth + 1, self.l[idx], ans)
            if self.r[idx] != -1 and y < yr:
                self._find(xl, xr, yl, yr, depth + 1, self.r[idx], ans)

    def find(self, xl, xr, yl, yr):
        ans = []
        idx = self.root
        depth = 0
        self._find(xl, xr, yl, yr, depth, idx, ans)
        return ans

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