class Node: """ Represents node of *R*-tree. Can be subclassed for custom metrics definition. """ __slots__ = 'index', 'box', 'children' def __init__(self, index: int, box: Box, children: Optional[Sequence['Node']]) -> None: self.index = index self.box = box self.children = children __repr__ = _generate_repr(__init__) @property def is_leaf(self) -> bool: """Checks whether the node is a leaf.""" return self.children is None @property def item(self) -> Item: """Returns underlying index with box.""" return self.index, self.box def distance_to_point(self, point: Point) -> Coordinate: """Calculates distance to given point.""" return _box.distance_to_point(self.box, point)
class Tree: """ Represents packed 2-dimensional segmental Hilbert *R*-tree. Reference: https://en.wikipedia.org/wiki/Hilbert_R-tree#Packed_Hilbert_R-trees """ __slots__ = '_context', '_max_children', '_root', '_segments' def __init__(self, segments: _Sequence[_Segment], *, max_children: int = 16, context: _Optional[_Context] = None) -> None: """ Initializes tree from segments. Time complexity: ``O(size * log size)`` Memory complexity: ``O(size)`` where ``size = len(segments)``. """ if context is None: context = _get_context() box_cls = context.box_cls self._context, self._max_children, self._root, self._segments = ( context, max_children, _create_root(segments, [box_cls(*((segment.start.x, segment.end.x) if segment.start.x < segment.end.x else (segment.end.x, segment.start.x)), *((segment.start.y, segment.end.y) if segment.start.y < segment.end.y else (segment.end.y, segment.start.y))) for segment in segments], max_children, context.merged_box, context.box_point_squared_distance, context.box_segment_squared_distance, context.segment_point_squared_distance, context.segments_squared_distance), segments) __repr__ = _generate_repr(__init__) @property def context(self) -> _Context: """ Returns context of the tree. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` """ return self._context @property def max_children(self) -> int: """ Returns maximum number of children in each node. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> tree.max_children == 16 True """ return self._max_children @property def segments(self) -> _Sequence[_Segment]: """ Returns underlying segments. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> tree.segments == segments True """ return self._segments def n_nearest_indices(self, n: int, segment: _Segment) -> _Sequence[int]: """ Searches for indices of segments in the tree the nearest to the given segment. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise where ``size = len(self.segments)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result indices. :param segment: input segment. :returns: indices of segments in the tree the nearest to the input segment. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.n_nearest_indices(2, Segment(Point(0, 0), Point(10, 0))) ... == [0, 1]) True >>> (tree.n_nearest_indices(10, Segment(Point(0, 0), Point(10, 0))) ... == range(len(segments))) True """ return ([index for index, _ in self._n_nearest_items(n, segment)] if n < len(self._segments) else range(len(self._segments))) def n_nearest_items(self, n: int, segment: _Segment) -> _Sequence[_Item]: """ Searches for indices with segments in the tree the nearest to the given segment. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(size)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(size)`` otherwise where ``size = len(self.segments)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result indices with segments. :param segment: input segment. :returns: indices with segments in the tree the nearest to the input segment. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.n_nearest_items(2, Segment(Point(0, 0), Point(10, 0))) ... == [(0, Segment(Point(0, 1), Point(1, 1))), ... (1, Segment(Point(0, 2), Point(2, 2)))]) True >>> (tree.n_nearest_items(10, Segment(Point(0, 0), Point(10, 0))) ... == list(enumerate(segments))) True """ return list(self._n_nearest_items(n, segment) if n < len(self._segments) else enumerate(self._segments)) def n_nearest_segments(self, n: int, segment: _Segment ) -> _Sequence[_Segment]: """ Searches for segments in the tree the nearest to the given segment. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise where ``size = len(self.segments)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result segments. :param segment: input segment. :returns: segments in the tree the nearest to the input segment. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.n_nearest_segments(2, Segment(Point(0, 0), Point(10, 0))) ... == [Segment(Point(0, 1), Point(1, 1)), ... Segment(Point(0, 2), Point(2, 2))]) True >>> (tree.n_nearest_segments(10, Segment(Point(0, 0), Point(10, 0))) ... == segments) True """ return ([segment for _, segment in self._n_nearest_items(n, segment)] if n < len(self._segments) else self._segments) def n_nearest_to_point_indices(self, n: int, point: _Point ) -> _Sequence[int]: """ Searches for indices of segments in the tree the nearest to the given point. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise where ``size = len(self.segments)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result indices. :param point: input point. :returns: indices of segments in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> tree.n_nearest_to_point_indices(2, Point(0, 0)) == [0, 1] True >>> (tree.n_nearest_to_point_indices(10, Point(0, 0)) ... == range(len(segments))) True """ return ([index for index, _ in self._n_nearest_to_point_items(n, point)] if n < len(self._segments) else range(len(self._segments))) def n_nearest_to_point_items(self, n: int, point: _Point ) -> _Sequence[_Item]: """ Searches for indices with segments in the tree the nearest to the given point. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(size)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(size)`` otherwise where ``size = len(self.segments)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result indices with segments. :param point: input point. :returns: indices with segments in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.n_nearest_to_point_items(2, Point(0, 0)) ... == [(0, Segment(Point(0, 1), Point(1, 1))), ... (1, Segment(Point(0, 2), Point(2, 2)))]) True >>> (tree.n_nearest_to_point_items(10, Point(0, 0)) ... == list(enumerate(segments))) True """ return list(self._n_nearest_to_point_items(n, point) if n < len(self._segments) else enumerate(self._segments)) def n_nearest_to_point_segments(self, n: int, point: _Point ) -> _Sequence[_Segment]: """ Searches for segments in the tree the nearest to the given point. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise where ``size = len(self.segments)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result segments. :param point: input point. :returns: segments in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.n_nearest_to_point_segments(2, Point(0, 0)) ... == [Segment(Point(0, 1), Point(1, 1)), ... Segment(Point(0, 2), Point(2, 2))]) True >>> tree.n_nearest_to_point_segments(10, Point(0, 0)) == segments True """ return ([segment for _, segment in self._n_nearest_to_point_items(n, point)] if n < len(self._segments) else self._segments) def nearest_index(self, segment: _Segment) -> int: """ Searches for index of segment in the tree the nearest to the given segment. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.segments)``, ``max_children = self.max_children``. :param segment: input segment. :returns: index of segment in the tree the nearest to the input segment. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> tree.nearest_index(Segment(Point(0, 0), Point(10, 0))) == 0 True """ result, _ = self.nearest_item(segment) return result def nearest_item(self, segment: _Segment) -> _Item: """ Searches for index with segment in the tree the nearest to the given segment. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.segments)``, ``max_children = self.max_children``. :param segment: input segment. :returns: index with segment in the tree the nearest to the input segment. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.nearest_item(Segment(Point(0, 0), Point(10, 0))) ... == (0, Segment(Point(0, 1), Point(1, 1)))) True """ queue = [(0, 0, self._root)] while queue: _, _, node = _heappop(queue) for child in node.children: _heappush(queue, (child.distance_to_segment(segment), child.index if child.is_leaf else -child.index - 1, child)) if queue and queue[0][1] >= 0: _, _, node = _heappop(queue) return node.item def nearest_segment(self, segment: _Segment) -> _Segment: """ Searches for segment in the tree the nearest to the given segment. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.segments)``, ``max_children = self.max_children``. :param segment: input segment. :returns: segment in the tree the nearest to the input segment. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.nearest_segment(Segment(Point(0, 0), Point(10, 0))) ... == Segment(Point(0, 1), Point(1, 1))) True """ _, result = self.nearest_item(segment) return result def nearest_to_point_index(self, point: _Point) -> int: """ Searches for index of segment in the tree the nearest to the given point. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.segments)``, ``max_children = self.max_children``. :param point: input point. :returns: index of segment in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> tree.nearest_to_point_index(Point(0, 0)) == 0 True """ result, _ = self.nearest_to_point_item(point) return result def nearest_to_point_item(self, point: _Point) -> _Item: """ Searches for index with segment in the tree the nearest to the given point. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.segments)``, ``max_children = self.max_children``. :param point: input point. :returns: index with segment in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.nearest_to_point_item(Point(0, 0)) ... == (0, Segment(Point(0, 1), Point(1, 1)))) True """ queue = [(0, 0, self._root)] while queue: _, _, node = _heappop(queue) for child in node.children: _heappush(queue, (child.distance_to_point(point), child.index if child.is_leaf else -child.index - 1, child)) if queue and queue[0][1] >= 0: _, _, node = _heappop(queue) return node.item def nearest_to_point_segment(self, point: _Point) -> _Segment: """ Searches for segment in the tree the nearest to the given point. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.segments)``, ``max_children = self.max_children``. :param point: input point. :returns: segment in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point, Segment = context.point_cls, context.segment_cls >>> segments = [Segment(Point(0, index), Point(index, index)) ... for index in range(1, 11)] >>> tree = Tree(segments) >>> (tree.nearest_to_point_segment(Point(0, 0)) ... == Segment(Point(0, 1), Point(1, 1))) True """ _, result = self.nearest_to_point_item(point) return result def _n_nearest_items(self, n: int, segment: _Segment) -> _Iterator[_Item]: queue = [(0, 0, self._root)] while n and queue: _, _, node = _heappop(queue) for child in node.children: _heappush(queue, (child.distance_to_segment(segment), child.index if child.is_leaf else -child.index - 1, child)) while n and queue and queue[0][1] >= 0: _, _, node = _heappop(queue) yield node.item n -= 1 def _n_nearest_to_point_items(self, n: int, point: _Point ) -> _Iterator[_Item]: queue = [(0, 0, self._root)] while n and queue: _, _, node = _heappop(queue) for child in node.children: _heappush(queue, (child.distance_to_point(point), child.index if child.is_leaf else -child.index - 1, child)) while n and queue and queue[0][1] >= 0: _, _, node = _heappop(queue) yield node.item n -= 1
class Tree: """ Represents packed 2-dimensional Hilbert *R*-tree. Reference: https://en.wikipedia.org/wiki/Hilbert_R-tree#Packed_Hilbert_R-trees """ __slots__ = '_boxes', '_context', '_max_children', '_root' def __init__(self, boxes: _Sequence[_Box], *, max_children: int = 16, context: _Optional[_Context] = None) -> None: """ Initializes tree from boxes. Time complexity: ``O(size * log size)`` Memory complexity: ``O(size)`` where ``size = len(boxes)``. """ if context is None: context = _get_context() self._boxes, self._context, self._max_children, self._root = ( boxes, context, max_children, _create_root(boxes, max_children, context.merged_box, context.box_point_squared_distance)) __repr__ = _generate_repr(__init__) @property def boxes(self) -> _Sequence[_Box]: """ Returns underlying boxes. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.boxes == boxes True """ return self._boxes @property def context(self) -> _Context: """ Returns context of the tree. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` """ return self._context @property def max_children(self) -> int: """ Returns maximum number of children in each node. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.max_children == 16 True """ return self._max_children def find_subsets(self, box: _Box) -> _List[_Box]: """ Searches for boxes that lie inside the given box. Time complexity: ``O(max_children * log size + hits_count)`` Memory complexity: ``O(max_children * log size + hits_count)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``, ``hits_count`` --- number of found boxes. :param box: input box. :returns: boxes that lie inside the input box. >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.find_subsets(Box(-1, 1, 0, 1)) == [Box(-1, 1, 0, 1)] True >>> (tree.find_subsets(Box(-2, 2, 0, 2)) ... == [Box(-1, 1, 0, 1), Box(-2, 2, 0, 2)]) True >>> (tree.find_subsets(Box(-3, 3, 0, 3)) ... == [Box(-1, 1, 0, 1), Box(-2, 2, 0, 2), Box(-3, 3, 0, 3)]) True """ return [box for _, box in self._find_subsets_items(box)] def find_subsets_indices(self, box: _Box) -> _List[int]: """ Searches for indices of boxes that lie inside the given box. Time complexity: ``O(max_children * log size + hits_count)`` Memory complexity: ``O(max_children * log size + hits_count)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``, ``hits_count`` --- number of found indices. :param box: input box. :returns: indices of boxes that lie inside the input box. >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.find_subsets_indices(Box(-1, 1, 0, 1)) == [0] True >>> tree.find_subsets_indices(Box(-2, 2, 0, 2)) == [0, 1] True >>> tree.find_subsets_indices(Box(-3, 3, 0, 3)) == [0, 1, 2] True """ return [index for index, _ in self._find_subsets_items(box)] def find_subsets_items(self, box: _Box) -> _List[_Item]: """ Searches for indices with boxes that lie inside the given box. Time complexity: ``O(max_children * log size + hits_count)`` Memory complexity: ``O(max_children * log size + hits_count)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``, ``hits_count`` --- number of found indices with boxes. :param box: input box. :returns: indices with boxes that lie inside the input box. >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> (tree.find_subsets_items(Box(-1, 1, 0, 1)) ... == [(0, Box(-1, 1, 0, 1))]) True >>> (tree.find_subsets_items(Box(-2, 2, 0, 2)) ... == [(0, Box(-1, 1, 0, 1)), (1, Box(-2, 2, 0, 2))]) True >>> (tree.find_subsets_items(Box(-3, 3, 0, 3)) ... == [(0, Box(-1, 1, 0, 1)), (1, Box(-2, 2, 0, 2)), ... (2, Box(-3, 3, 0, 3))]) True """ return list(self._find_subsets_items(box)) def find_supersets(self, box: _Box) -> _List[_Box]: """ Searches for boxes that contain the given box. Time complexity: ``O(max_children * log size + hits_count)`` Memory complexity: ``O(max_children * log size + hits_count)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``, ``hits_count`` --- number of found boxes. :param box: input box. :returns: boxes that contain the input box. >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.find_supersets(Box(-10, 10, 0, 10)) == [Box(-10, 10, 0, 10)] True >>> (tree.find_supersets(Box(-9, 9, 0, 9)) ... == [Box(-9, 9, 0, 9), Box(-10, 10, 0, 10)]) True >>> (tree.find_supersets(Box(-8, 8, 0, 8)) ... == [Box(-8, 8, 0, 8), Box(-9, 9, 0, 9), Box(-10, 10, 0, 10)]) True """ return [box for _, box in self._find_supersets_items(box)] def find_supersets_indices(self, box: _Box) -> _List[int]: """ Searches for indices of boxes that contain the given box. Time complexity: ``O(max_children * log size + hits_count)`` Memory complexity: ``O(max_children * log size + hits_count)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``, ``hits_count`` --- number of found indices. :param box: input box. :returns: indices of boxes that contain the input box. >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.find_supersets_indices(Box(-10, 10, 0, 10)) == [9] True >>> tree.find_supersets_indices(Box(-9, 9, 0, 9)) == [8, 9] True >>> tree.find_supersets_indices(Box(-8, 8, 0, 8)) == [7, 8, 9] True """ return [index for index, _ in self._find_supersets_items(box)] def find_supersets_items(self, box: _Box) -> _List[_Item]: """ Searches for indices with boxes that contain the given box. Time complexity: ``O(max_children * log size + hits_count)`` Memory complexity: ``O(max_children * log size + hits_count)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``, ``hits_count`` --- number of found indices with boxes. :param box: input box. :returns: indices with boxes that contain the input box. >>> from ground.base import get_context >>> context = get_context() >>> Box = context.box_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> (tree.find_supersets_items(Box(-10, 10, 0, 10)) ... == [(9, Box(-10, 10, 0, 10))]) True >>> (tree.find_supersets_items(Box(-9, 9, 0, 9)) ... == [(8, Box(-9, 9, 0, 9)), (9, Box(-10, 10, 0, 10))]) True >>> (tree.find_supersets_items(Box(-8, 8, 0, 8)) ... == [(7, Box(-8, 8, 0, 8)), (8, Box(-9, 9, 0, 9)), ... (9, Box(-10, 10, 0, 10))]) True """ return list(self._find_supersets_items(box)) def _find_subsets_items(self, box: _Box) -> _Iterator[_Item]: yield from (enumerate(self._boxes) if _box.is_subset_of( self._root.box, box) else _find_node_box_subsets_items( self._root, box)) def _find_supersets_items(self, box: _Box) -> _Iterator[_Item]: yield from _find_node_box_supersets_items(self._root, box) def n_nearest_indices(self, n: int, point: _Point) -> _Sequence[int]: """ Searches for indices of boxes in the tree the nearest to the given point. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise where ``size = len(self.boxes)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result indices. :param point: input point. :returns: indices of boxes in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.n_nearest_indices(2, Point(0, 0)) == [9, 8] True >>> (tree.n_nearest_indices(len(boxes), Point(0, 0)) ... == range(len(boxes))) True """ return ([index for index, _ in self._n_nearest_items(n, point)] if n < len(self._boxes) else range(len(self._boxes))) def n_nearest_boxes(self, n: int, point: _Point) -> _Sequence[_Box]: """ Searches for boxes in the tree the nearest to the given point. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(1)`` otherwise where ``size = len(self.boxes)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result boxes. :param point: input point. :returns: boxes in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> (tree.n_nearest_boxes(2, Point(0, 0)) ... == [Box(-10, 10, 0, 10), Box(-9, 9, 0, 9)]) True >>> tree.n_nearest_boxes(len(boxes), Point(0, 0)) == boxes True """ return ([box for _, box in self._n_nearest_items(n, point)] if n < len(self._boxes) else self._boxes) def n_nearest_items(self, n: int, point: _Point) -> _Sequence[_Item]: """ Searches for indices with boxes in the tree the nearest to the given point. Time complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(size)`` otherwise Memory complexity: ``O(n * max_children * log size)`` if ``n < size``, ``O(size)`` otherwise where ``size = len(self.boxes)``, ``max_children = self.max_children``. :param n: positive upper bound for number of result indices with boxes. :param point: input point. :returns: indices with boxes in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> (tree.n_nearest_items(2, Point(0, 0)) ... == [(9, Box(-10, 10, 0, 10)), (8, Box(-9, 9, 0, 9))]) True >>> (tree.n_nearest_items(len(boxes), Point(0, 0)) ... == list(enumerate(boxes))) True """ return list( self._n_nearest_items(n, point) if n < len(self._boxes) else enumerate(self._boxes)) def nearest_index(self, point: _Point) -> int: """ Searches for index of box in the tree the nearest to the given point. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``. :param point: input point. :returns: index of box in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.nearest_index(Point(0, 0)) == 9 True """ result, _ = self.nearest_item(point) return result def nearest_box(self, point: _Point) -> _Box: """ Searches for box in the tree the nearest to the given point. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``. :param point: input point. :returns: box in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.nearest_box(Point(0, 0)) == Box(-10, 10, 0, 10) True """ _, result = self.nearest_item(point) return result def nearest_item(self, point: _Point) -> _Item: """ Searches for index with box in the tree the nearest to the given point. Time complexity: ``O(max_children * log size)`` Memory complexity: ``O(max_children * log size)`` where ``size = len(self.boxes)``, ``max_children = self.max_children``. :param point: input point. :returns: index with box in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> boxes = [Box(-index, index, 0, index) for index in range(1, 11)] >>> tree = Tree(boxes) >>> tree.nearest_item(Point(0, 0)) == (9, Box(-10, 10, 0, 10)) True >>> tree.nearest_item(Point(-10, 0)) == (9, Box(-10, 10, 0, 10)) True >>> tree.nearest_item(Point(-10, 10)) == (9, Box(-10, 10, 0, 10)) True >>> tree.nearest_item(Point(10, 0)) == (9, Box(-10, 10, 0, 10)) True >>> tree.nearest_item(Point(10, 10)) == (9, Box(-10, 10, 0, 10)) True """ queue = [(0, 0, self._root)] while queue: _, _, node = _heappop(queue) for child in node.children: _heappush(queue, (child.distance_to_point(point), -child.index - 1 if child.is_leaf else child.index, child)) if queue and queue[0][1] < 0: _, _, node = _heappop(queue) return node.item def _n_nearest_items(self, n: int, point: _Point) -> _Iterator[_Item]: queue = [(0, 0, self._root)] while n and queue: _, _, node = _heappop(queue) for child in node.children: _heappush(queue, (child.distance_to_point(point), -child.index - 1 if child.is_leaf else child.index, child)) while n and queue and queue[0][1] < 0: _, _, node = _heappop(queue) yield node.item n -= 1
class Context: __slots__ = ('_box_cls', '_centroidal', '_contour_cls', '_incircle', '_linear', '_mode', '_multipoint_cls', '_multipolygon_cls', '_multisegment_cls', '_point_cls', '_polygon_cls', '_segment_cls', '_vector') def __init__(self, *, box_cls: _Type[_hints.Box] = _geometries.Box, contour_cls: _Type[_hints.Contour] = _geometries.Contour, multipoint_cls: _Type[ _hints.Multipoint] = _geometries.Multipoint, multipolygon_cls: _Type[ _hints.Multipolygon] = _geometries.Multipolygon, multisegment_cls: _Type[ _hints.Multisegment] = _geometries.Multisegment, point_cls: _Type[_hints.Point] = _geometries.Point, polygon_cls: _Type[_hints.Polygon] = _geometries.Polygon, segment_cls: _Type[_hints.Segment] = _geometries.Segment, mode: Mode = Mode.EXACT) -> None: self._box_cls = box_cls self._contour_cls = contour_cls self._multipoint_cls = multipoint_cls self._multipolygon_cls = multipolygon_cls self._multisegment_cls = multisegment_cls self._point_cls = point_cls self._polygon_cls = polygon_cls self._segment_cls = segment_cls self._mode = mode self._centroidal, self._incircle, self._linear, self._vector = ( (_centroidal.exact_context, _incircle.exact_context, _linear.exact_context, _vector.exact_context) if mode is Mode.EXACT else ((_centroidal.plain_context, _incircle.plain_context, _linear.plain_context, _vector.plain_context) if mode is Mode.PLAIN else (_centroidal.robust_context, _incircle.robust_context, _linear.exact_context, _vector.robust_context))) __repr__ = _generate_repr(__init__) @property def box_cls(self) -> _Type[_hints.Box]: return self._box_cls @property def contour_cls(self) -> _Type[_hints.Contour]: return self._contour_cls @property def cross_product(self) -> _QuaternaryFunction: return self._vector.cross_product @property def dot_product(self) -> _QuaternaryFunction: return self._vector.dot_product @property def mode(self) -> Mode: return self._mode @property def multipoint_cls(self) -> _Type[_hints.Multipoint]: return self._multipoint_cls @property def multipolygon_cls(self) -> _Type[_hints.Multipolygon]: return self._multipolygon_cls @property def multisegment_cls(self) -> _Type[_hints.Multisegment]: return self._multisegment_cls @property def point_cls(self) -> _Type[_hints.Point]: return self._point_cls @property def point_point_point_incircle_test(self) -> _QuaternaryFunction: return self._incircle.point_point_point_test @property def polygon_cls(self) -> _Type[_hints.Polygon]: return self._polygon_cls @property def segment_cls(self) -> _Type[_hints.Segment]: return self._segment_cls def angle_kind(self, vertex: _hints.Point, first_ray_point: _hints.Point, second_ray_point: _hints.Point) -> Kind: return _angular.kind(self.dot_product, vertex, first_ray_point, second_ray_point) def angle_orientation(self, vertex: _hints.Point, first_ray_point: _hints.Point, second_ray_point: _hints.Point) -> Orientation: return _angular.orientation(self.cross_product, vertex, first_ray_point, second_ray_point) def contour_centroid(self, vertices: _Sequence[_hints.Point]) -> _hints.Point: """ Constructs centroid of a contour given its vertices. Time complexity: ``O(len(vertices))`` Memory complexity: ``O(1)`` >>> context = get_context() >>> Point = context.point_cls >>> context.contour_centroid([Point(0, 0), Point(2, 0), Point(2, 2), ... Point(0, 2)]) == Point(1, 1) True """ return self._centroidal.contour_centroid(self.point_cls, vertices) def merged_box(self, first_box: _hints.Box, second_box: _hints.Box) -> _hints.Box: return _boxed.merge(self.box_cls, first_box, second_box) def multipoint_centroid(self, points: _Sequence[_hints.Point]) -> _hints.Point: """ Constructs centroid of a multipoint given its points. Time complexity: ``O(len(points))`` Memory complexity: ``O(1)`` >>> context = get_context() >>> Point = context.point_cls >>> context.multipoint_centroid([Point(0, 0), Point(2, 0), Point(2, 2), ... Point(0, 2)]) == Point(1, 1) True """ return self._centroidal.multipoint_centroid(self.point_cls, points) def points_convex_hull( self, points: _Sequence[_hints.Point]) -> _Sequence[_hints.Point]: """ Constructs convex hull of points. Time complexity: ``O(points_count * log(points_count))`` Memory complexity: ``O(points_count)`` where ``points_count = len(points)``. >>> context = get_context() >>> Point = context.point_cls >>> (context.points_convex_hull([Point(0, 0), Point(2, 0), Point(2, 2), ... Point(0, 2)]) ... == [Point(0, 0), Point(2, 0), Point(2, 2), Point(0, 2)]) True """ return _discrete.to_convex_hull(self.angle_orientation, points) def segment_contains_point(self, start: _hints.Point, end: _hints.Point, point: _hints.Point) -> bool: """ Checks if a segment given by its endpoints contains given point. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` >>> context = get_context() >>> Point = context.point_cls >>> context.segment_contains_point(Point(0, 0), Point(2, 0), ... Point(0, 0)) True >>> context.segment_contains_point(Point(0, 0), Point(2, 0), ... Point(0, 2)) False >>> context.segment_contains_point(Point(0, 0), Point(2, 0), ... Point(1, 0)) True >>> context.segment_contains_point(Point(0, 0), Point(2, 0), ... Point(1, 1)) False >>> context.segment_contains_point(Point(0, 0), Point(2, 0), ... Point(2, 0)) True >>> context.segment_contains_point(Point(0, 0), Point(2, 0), ... Point(3, 0)) False """ return self._linear.containment_checker(self.cross_product, start, end, point) def segments_intersection(self, first_start: _hints.Point, first_end: _hints.Point, second_start: _hints.Point, second_end: _hints.Point) -> _hints.Point: return self._linear.intersector(self.cross_product, self.point_cls, first_start, first_end, second_start, second_end) def segments_relation(self, test_start: _hints.Point, test_end: _hints.Point, goal_start: _hints.Point, goal_end: _hints.Point) -> Relation: return self._linear.relater(self.cross_product, test_start, test_end, goal_start, goal_end)
class Tree: """ Represents `k`-dimensional (aka *kd*) tree. Reference: https://en.wikipedia.org/wiki/K-d_tree """ __slots__ = '_context', '_points', '_root' def __init__(self, points: _Sequence[_Point], *, context: _Optional[_Context] = None) -> None: """ Initializes tree from points. Time complexity: ``O(dimension * size * log size)`` Memory complexity: ``O(dimension * size)`` where ``dimension = len(points[0])``, ``size = len(points)``. """ if context is None: context = _get_context() self._context, self._points, self._root = ( context, points, _create_node(range(len(points)), points, False, context.points_squared_distance)) __repr__ = _generate_repr(__init__) @property def context(self) -> _Context: """ Returns context of the tree. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` """ return self._context @property def points(self) -> _Sequence[_Point]: """ Returns underlying points. Time complexity: ``O(1)`` Memory complexity: ``O(1)`` >>> from ground.base import get_context >>> context = get_context() >>> Point = context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.points == points True """ return self._points def n_nearest_indices(self, n: int, point: _Point) -> _Sequence[int]: """ Searches for indices of points in the tree that are the nearest to the given point. Time complexity: ``O(min(n, size) * log size)`` Memory complexity: ``O(min(n, size) * log size)`` where ``size = len(self.points)``. Reference: https://en.wikipedia.org/wiki/K-d_tree#Nearest_neighbour_search :param n: positive upper bound for number of result indices. :param point: input point. :returns: indices of points in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point = context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.n_nearest_indices(2, Point(0, 0)) == [2, 3] True >>> (tree.n_nearest_indices(len(points), Point(0, 0)) ... == range(len(points))) True """ return ([index for index, _ in self._n_nearest_items(n, point)] if n < len(self._points) else range(len(self._points))) def n_nearest_points(self, n: int, point: _Point) -> _Sequence[_Point]: """ Searches for points in the tree the nearest to the given point. Time complexity: ``O(min(n, size) * log size)`` Memory complexity: ``O(min(n, size) * log size)`` where ``size = len(self.points)``. Reference: https://en.wikipedia.org/wiki/K-d_tree#Nearest_neighbour_search :param n: positive upper bound for number of result points. :param point: input point. :returns: points in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point = context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> (tree.n_nearest_points(2, Point(0, 0)) ... == [Point(-3, 2), Point(-2, 3)]) True >>> tree.n_nearest_points(len(points), Point(0, 0)) == points True """ return ([point for _, point in self._n_nearest_items(n, point)] if n < len(self._points) else self._points) def n_nearest_items(self, n: int, point: _Point) -> _Sequence[_Item]: """ Searches for indices with points in the tree that are the nearest to the given point. Time complexity: ``O(min(n, size) * log size)`` Memory complexity: ``O(min(n, size) * log size)`` where ``size = len(self.points)``. Reference: https://en.wikipedia.org/wiki/K-d_tree#Nearest_neighbour_search :param n: positive upper bound for number of result indices. :param point: input point. :returns: indices with points in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point = context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> (tree.n_nearest_items(2, Point(0, 0)) ... == [(2, Point(-3, 2)), (3, Point(-2, 3))]) True >>> (tree.n_nearest_items(len(points), Point(0, 0)) ... == list(enumerate(points))) True """ return (self._n_nearest_items(n, point) if n < len(self._points) else list(enumerate(self._points))) def _n_nearest_items(self, n: int, point: _Point) -> _List[_Item]: candidates = [] # type: _List[_Tuple[_Scalar, _Item]] queue = [self._root] push, pop = queue.append, queue.pop while queue: node = pop() # type: _Node distance_to_point = node.distance_to_point(point) candidate = -distance_to_point, node.item if len(candidates) < n: _heappush(candidates, candidate) elif distance_to_point < -candidates[0][0]: _heapreplace(candidates, candidate) coordinate = node.projector(point) point_is_on_the_left = coordinate < node.projection if point_is_on_the_left: if node.left is not _NIL: push(node.left) elif node.right is not _NIL: push(node.right) if (len(candidates) < n or (node.distance_to_coordinate(coordinate) < -candidates[0][0])): if point_is_on_the_left: if node.right is not _NIL: push(node.right) elif node.left is not _NIL: push(node.left) return [item for _, item in candidates] def nearest_index(self, point: _Point) -> int: """ Searches for index of a point in the tree that is the nearest to the given point. Time complexity: ``O(log size)`` Memory complexity: ``O(log size)`` where ``size = len(self.points)``. Reference: https://en.wikipedia.org/wiki/K-d_tree#Nearest_neighbour_search :param point: input point. :returns: index of a point in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point = context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.nearest_index(Point(0, 0)) == 2 True >>> tree.nearest_index(Point(-3, 2)) == 2 True """ result, _ = self.nearest_item(point) return result def nearest_point(self, point: _Point) -> _Point: """ Searches for point in the tree that is the nearest to the given point. Time complexity: ``O(log size)`` Memory complexity: ``O(log size)`` where ``size = len(self.points)``. Reference: https://en.wikipedia.org/wiki/K-d_tree#Nearest_neighbour_search :param point: input point. :returns: point in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point = context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.nearest_point(Point(0, 0)) == Point(-3, 2) True >>> tree.nearest_point(Point(-3, 2)) == Point(-3, 2) True """ _, result = self.nearest_item(point) return result def nearest_item(self, point: _Point) -> _Item: """ Searches for index with point in the tree that is the nearest to the given point. Time complexity: ``O(log size)`` Memory complexity: ``O(log size)`` where ``size = len(self.points)``. Reference: https://en.wikipedia.org/wiki/K-d_tree#Nearest_neighbour_search :param point: input point. :returns: index with point in the tree the nearest to the input point. >>> from ground.base import get_context >>> context = get_context() >>> Point = context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.nearest_item(Point(0, 0)) == (2, Point(-3, 2)) True >>> tree.nearest_item(Point(-3, 2)) == (2, Point(-3, 2)) True """ node = self._root result, min_distance = node.item, node.distance_to_point(point) queue = [node] push, pop = queue.append, queue.pop while queue: node = pop() # type: _Node distance_to_point = node.distance_to_point(point) if distance_to_point < min_distance: result, min_distance = node.item, distance_to_point coordinate = node.projector(point) point_is_on_the_left = coordinate < node.projection if point_is_on_the_left: if node.left is not _NIL: push(node.left) elif node.right is not _NIL: push(node.right) if node.distance_to_coordinate(coordinate) < min_distance: if point_is_on_the_left: if node.right is not _NIL: push(node.right) elif node.left is not _NIL: push(node.left) return result def find_box_indices(self, box: _Box) -> _List[int]: """ Searches for indices of points that lie inside the given box. Time complexity: ``O(dimension * size ** (1 - 1 / dimension) + hits_count)`` Memory complexity: ``O(dimension * size ** (1 - 1 / dimension) + hits_count)`` where ``dimension = len(self.points[0])``, ``size = len(self.points)``, ``hits_count`` --- number of found indices. Reference: https://en.wikipedia.org/wiki/K-d_tree#Range_search :param box: box to search in. :returns: indices of points that lie inside the box. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.find_box_indices(Box(-3, 3, 0, 1)) == [] True >>> tree.find_box_indices(Box(-3, 3, 0, 2)) == [2] True >>> tree.find_box_indices(Box(-3, 3, 0, 3)) == [2, 3] True """ return [index for index, _ in self._find_box_items(box)] def find_box_points(self, box: _Box) -> _List[_Point]: """ Searches for points that lie inside the given box. Time complexity: ``O(dimension * size ** (1 - 1 / dimension) + hits_count)`` Memory complexity: ``O(dimension * size ** (1 - 1 / dimension) + hits_count)`` where ``dimension = len(self.points[0])``, ``size = len(self.points)``, ``hits_count`` --- number of found points. Reference: https://en.wikipedia.org/wiki/K-d_tree#Range_search :param box: box to search in. :returns: points that lie inside the box. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.find_box_points(Box(-3, 3, 0, 1)) == [] True >>> tree.find_box_points(Box(-3, 3, 0, 2)) == [Point(-3, 2)] True >>> (tree.find_box_points(Box(-3, 3, 0, 3)) ... == [Point(-3, 2), Point(-2, 3)]) True """ return [point for _, point in self._find_box_items(box)] def find_box_items(self, box: _Box) -> _List[_Item]: """ Searches for indices with points in the tree that lie inside the given box. Time complexity: ``O(dimension * size ** (1 - 1 / dimension) + hits_count)`` Memory complexity: ``O(dimension * size ** (1 - 1 / dimension) + hits_count)`` where ``dimension = len(self.points[0])``, ``size = len(self.points)``, ``hits_count`` --- number of found indices with points. Reference: https://en.wikipedia.org/wiki/K-d_tree#Range_search :param box: box to search in. :returns: indices with points in the tree that lie inside the box. >>> from ground.base import get_context >>> context = get_context() >>> Box, Point = context.box_cls, context.point_cls >>> points = list(map(Point, range(-5, 6), range(10))) >>> tree = Tree(points) >>> tree.find_box_items(Box(-3, 3, 0, 1)) == [] True >>> tree.find_box_items(Box(-3, 3, 0, 2)) == [(2, Point(-3, 2))] True >>> (tree.find_box_items(Box(-3, 3, 0, 3)) ... == [(2, Point(-3, 2)), (3, Point(-2, 3))]) True """ return list(self._find_box_items(box)) def _find_box_items(self, box: _Box) -> _Iterator[_Item]: queue = [self._root] push, pop = queue.append, queue.pop while queue: node = pop() # type: _Node if _box.contains_point(box, node.point): yield node.item min_coordinate, max_coordinate = ((box.min_y, box.max_y) if node.is_y_axis else (box.min_x, box.max_x)) coordinate = node.projector(node.point) if node.left is not _NIL and min_coordinate <= coordinate: push(node.left) if node.right is not _NIL and coordinate <= max_coordinate: push(node.right)