kdtree.py

"""!
 
@brief Data Structure: KD-Tree
@details Implementation based on paper @cite book::the_design_and_analysis.
 
@authors Andrei Novikov (pyclustering@yandex.ru)
@date 2014-2020
@copyright BSD-3-Clause
 
"""
 
 
import numpy
import matplotlib.pyplot as plt
 
from pyclustering.utils import euclidean_distance_square, find_left_element
 
 
class kdtree_visualizer:
    """!
    @brief KD-tree visualizer that provides service to display graphical representation of the tree using
            `matplotlib` library.
    @details The visualizer is able to visualize 2D KD-trees only.
 
    There is an example how to visualize balanced KD-tree for `TwoDiamonds` sample using `kdtree_visualizer`:
    @code
        from pyclustering.container.kdtree import kdtree_balanced, kdtree_visualizer
        from pyclustering.utils import read_sample
        from pyclustering.samples.definitions import FCPS_SAMPLES
 
        sample = read_sample(FCPS_SAMPLES.SAMPLE_TWO_DIAMONDS)
        tree_instance = kdtree_balanced(sample)
 
        kdtree_visualizer(tree_instance).visualize()
    @endcode
 
    Output result of the example above (balanced tree) - figure 1:
    @image html kd_tree_unbalanced_two_diamonds.png "Fig. 1. Balanced KD-tree for sample 'TwoDiamonds'."
 
    There is one more example to demonstrate unbalanced KD-tree. `kdtree` class is child class of `kdtree_balanced`
    that allows to add points step by step and thus an unbalanced KD-tree can be built.
    @code
        from pyclustering.container.kdtree import kdtree, kdtree_visualizer
        from pyclustering.utils import read_sample
        from pyclustering.samples.definitions import FCPS_SAMPLES
 
        sample = read_sample(FCPS_SAMPLES.SAMPLE_TWO_DIAMONDS)
        tree_instance = kdtree()    # Do not use sample in constructor to avoid building of balanced tree.
 
        # Fill KD-tree step by step to obtain unbalanced tree.
        for point in sample:
            tree_instance.insert(point)
 
        kdtree_visualizer(tree_instance).visualize()
    @endcode
 
    Output result of the example above (unbalanced tree) - figure 2:
    @image html kd_tree_unbalanced_two_diamonds.png "Fig. 2. Unbalanced KD-tree for sample 'TwoDiamonds'."
 
    """
 
    def __init__(self, kdtree_instance):
        """!
        @brief Initialize KD-tree visualizer.
 
        @param[in] kdtree_instance (kdtree): Instance of a KD-tree that should be visualized.
 
        """
        self.__tree = kdtree_instance
        self.__colors = ['blue', 'red', 'green']
        self.__verify()
 
    def visualize(self):
        """!
        @brief Visualize KD-tree using plot in 2-dimensional data-space.
 
        """
        node = self.__tree.get_root()
        min, max = self.__get_limits()
 
        figure = plt.figure(111)
        ax = figure.add_subplot(111)
 
        self.__draw_node(ax, node, min, max)
 
        ax.set_xlim(min[0], max[0])
        ax.set_ylim(min[1], max[1])
        plt.show()
 
    def __draw_node(self, ax, node, min, max):
        self.__draw_split_line(ax, node, min, max)
 
        if node.left is not None:
            rborder = max[:]
            rborder[node.disc] = node.data[node.disc]
            self.__draw_node(ax, node.left, min, rborder)
 
        if node.right is not None:
            lborder = min[:]
            lborder[node.disc] = node.data[node.disc]
            self.__draw_node(ax, node.right, lborder, max)
 
    def __draw_split_line(self, ax, node, min, max):
        max_coord = max[:]
        min_coord = min[:]
 
        dimension = len(min)
        for d in range(dimension):
            if d == node.disc:
                max_coord[d] = node.data[d]
                min_coord[d] = node.data[d]
 
        if dimension == 2:
            ax.plot(node.data[0], node.data[1], color='black', marker='.', markersize=6)
            ax.plot([min_coord[0], max_coord[0]], [min_coord[1], max_coord[1]], color=self.__colors[node.disc],
                    linestyle='-', linewidth=1)
 
    def __get_limits(self):
        dimension = len(self.__tree.get_root().data)
        nodes = self.__get_all_nodes()
 
        max, min = [float('-inf')] * dimension, [float('+inf')] * dimension
 
        for node in nodes:
            for d in range(dimension):
                if max[d] < node.data[d]:
                    max[d] = node.data[d]
 
                if min[d] > node.data[d]:
                    min[d] = node.data[d]
 
        return min, max
 
    def __get_all_nodes(self):
        nodes = []
 
        next_level = [self.__tree.get_root()]
 
        while len(next_level) != 0:
            cur_level = next_level
            nodes += next_level
            next_level = []
 
            for cur_node in cur_level:
                children = cur_node.get_children()
                if children is not None:
                    next_level += children
 
        return nodes
 
    def __verify(self):
        root = self.__tree.get_root()
        if root is None:
            raise ValueError("KD-Tree is empty - nothing to visualize.")
 
        dimension = len(root.data)
        if dimension != 2:
            raise NotImplementedError("KD-Tree data has '%d' dimension - only KD-tree with 2D data can be visualized."
                                      % dimension)
 
 
 
class node:
    """!
    @brief Represents a node in a KD-Tree.
    @details The KD-Tree node contains point's coordinates, discriminator, payload and pointers to parent and children.
 
    @see kdtree_balanced
    @see kdtree
 
    """
 
    def __init__(self, data=None, payload=None, left=None, right=None, disc=None, parent=None):
        """!
        @brief Creates KD-tree node.
 
        @param[in] data (list): Data point that is presented as list of coordinates.
        @param[in] payload (any): Payload of node (pointer to essence that is attached to this node).
        @param[in] left (node): Node of KD-Tree that represents left successor.
        @param[in] right (node): Node of KD-Tree that represents right successor.
        @param[in] disc (uint): Index of dimension of that node.
        @param[in] parent (node): Node of KD-Tree that represents parent.
 
        """
 
        
        self.data = data
 
        
        self.payload = payload
 
        
        self.left = left
 
        
        self.right = right
 
        
        self.disc = disc
 
        
        self.parent = parent
 
    def __repr__(self):
        """!
        @return (string) Default representation of the node.
 
        """
        left = None
        right = None
 
        if self.left is not None:
            left = self.left.data
 
        if self.right is not None:
            right = self.right.data
 
        return "(%s: [L:'%s', R:'%s'])" % (self.data, left, right)
 
    def __str__(self):
        """!
        @return (string) String representation of the node.
 
        """
        return self.__repr__()
 
    def get_children(self):
        """!
        @brief Returns list of not `None` children of the node.
 
        @return (list) list of not `None` children of the node; if the node does not have children
                        then `None` is returned.
 
        """
 
        if self.left is not None:
            yield self.left
        if self.right is not None:
            yield self.right
 
 
 
class kdtree_balanced:
    """!
    @brief Represents balanced static KD-tree that does not provide services to add and remove nodes after
            initialization.
    @details In the term KD tree, k denotes the dimensionality of the space being represented. Each data point is
              represented as a node in the k-d tree in the form of a record of type node.
 
    There is an example how to create KD-tree:
    @code
        from pyclustering.container.kdtree import kdtree_balanced, kdtree_visualizer
        from pyclustering.utils import read_sample
        from pyclustering.samples.definitions import FCPS_SAMPLES
 
        sample = read_sample(FCPS_SAMPLES.SAMPLE_LSUN)
        tree_instance = kdtree_balanced(sample)
 
        kdtree_visualizer(tree_instance).visualize()
    @endcode
 
    Output result of the example above - figure 1.
    @image html kd_tree_balanced_lsun.png "Fig. 1. Balanced KD-tree for sample 'Lsun'."
 
    @see kdtree
 
    """
 
    def __init__(self, points, payloads=None):
        """!
        @brief Initializes balanced static KD-tree.
 
        @param[in] points (array_like): Points that should be used to build KD-tree.
        @param[in] payloads (array_like): Payload of each point in `points`.
 
        """
 
        if points is None:
            self._length = 0
            self._dimension = 0
            self._point_comparator = None
            self._root = None
            return
 
        self._dimension = len(points[0])
        self._point_comparator = self._create_point_comparator(type(points[0]))
        self._length = 0
 
        nodes = []
        for i in range(len(points)):
            payload = None
            if payloads is not None:
                payload = payloads[i]
 
            nodes.append(node(points[i], payload, None, None, -1, None))
 
        self._root = self.__create_tree(nodes, None, 0)
 
    def __len__(self):
        """!
        @brief Returns amount of nodes in the KD-tree.
 
        @return (uint) Amount of nodes in the KD-tree.
 
        """
        return self._length
 
    def get_root(self):
        """!
        @brief Returns root of the tree.
 
        @return (node) The root of the tree.
 
        """
        return self._root
 
    def __create_tree(self, nodes, parent, depth):
        """!
        @brief Creates balanced sub-tree using elements from list `nodes`.
 
        @param[in] nodes (list): List of KD-tree nodes.
        @param[in] parent (node): Parent node that is used as a root to build the sub-tree.
        @param[in] depth (uint): Depth of the tree that where children of the `parent` should be placed.
 
        @return (node) Returns a node that is a root of the built sub-tree.
 
        """
        if len(nodes) == 0:
            return None
 
        discriminator = depth % self._dimension
 
        nodes.sort(key=lambda n: n.data[discriminator])
        median = len(nodes) // 2
 
        # Elements could be the same around the median, but all elements that are >= to the current should
        # be at the right side.
        # TODO: optimize by binary search - no need to use O(n)
        median = find_left_element(nodes, median, lambda n1, n2: n1.data[discriminator] < n2.data[discriminator])
        # while median - 1 >= 0 and \
        #         nodes[median].data[discriminator] == nodes[median - 1].data[discriminator]:
        #     median -= 1
 
        new_node = nodes[median]
        new_node.disc = discriminator
        new_node.parent = parent
        new_node.left = self.__create_tree(nodes[:median], new_node, depth + 1)
        new_node.right = self.__create_tree(nodes[median + 1:], new_node, depth + 1)
 
        self._length += 1
        return new_node
 
    def _create_point_comparator(self, type_point):
        """!
        @brief Create point comparator.
        @details In case of numpy.array specific comparator is required.
 
        @param[in] type_point (data_type): Type of point that is stored in KD-node.
 
        @return (callable) Callable point comparator to compare to points.
 
        """
        if type_point == numpy.ndarray:
            return lambda obj1, obj2: numpy.array_equal(obj1, obj2)
 
        return lambda obj1, obj2: obj1 == obj2
 
    def _find_node_by_rule(self, point, search_rule, cur_node):
        """!
        @brief Search node that satisfy to parameters in search rule.
        @details If node with specified parameters does not exist then None will be returned,
                  otherwise required node will be returned.
 
        @param[in] point (list): Coordinates of the point whose node should be found.
        @param[in] search_rule (lambda): Rule that is called to check whether node satisfies to search parameter.
        @param[in] cur_node (node): Node from which search should be started.
 
        @return (node) Node if it satisfies to input parameters, otherwise it return None.
 
        """
 
        if cur_node is None:
            cur_node = self._root
 
        while cur_node:
            if cur_node.data[cur_node.disc] <= point[cur_node.disc]:
                # no need to check each node, only when it may satisfy the condition
                if search_rule(cur_node):  # compare point with point in the current node
                    return cur_node
 
                cur_node = cur_node.right
            else:
                cur_node = cur_node.left
 
        return None
 
    def find_node_with_payload(self, point, point_payload, cur_node=None):
        """!
        @brief Find node with specified coordinates and payload.
        @details If node with specified parameters does not exist then None will be returned,
                  otherwise required node will be returned.
 
        @param[in] point (list): Coordinates of the point whose node should be found.
        @param[in] point_payload (any): Payload of the node that is searched in the tree.
        @param[in] cur_node (node): Node from which search should be started.
 
        @return (node) Node if it satisfies to input parameters, otherwise it return None.
 
        """
 
        rule_search = lambda node, point=point, payload=point_payload: self._point_comparator(node.data, point) and \
                                                                       node.payload == payload
        return self._find_node_by_rule(point, rule_search, cur_node)
 
    def find_node(self, point, cur_node=None):
        """!
        @brief Find node with coordinates that are defined by specified point.
        @details If node with specified parameters does not exist then None will be returned,
                  otherwise required node will be returned.
 
        @param[in] point (list): Coordinates of the point whose node should be found.
        @param[in] cur_node (node): Node from which search should be started.
 
        @return (node) Node if it satisfies to input parameters, otherwise it return None.
 
        """
 
        rule_search = lambda node, point=point: self._point_comparator(node.data, point)
        return self._find_node_by_rule(point, rule_search, cur_node)
 
    def find_nearest_dist_node(self, point, distance, retdistance=False):
        """!
        @brief Find nearest neighbor in area with radius = distance.
 
        @param[in] point (list): Maximum distance where neighbors are searched.
        @param[in] distance (double): Maximum distance where neighbors are searched.
        @param[in] retdistance (bool): If True - returns neighbors with distances to them, otherwise only neighbors
                    is returned.
 
        @return (node|list) Nearest neighbor if 'retdistance' is False and list with two elements [node, distance]
                 if 'retdistance' is True, where the first element is pointer to node and the second element is
                 distance to it.
 
        """
 
        best_nodes = self.find_nearest_dist_nodes(point, distance)
 
        if len(best_nodes) == 0:
            return None
 
        nearest = min(best_nodes, key=lambda item: item[0])
 
        if retdistance is True:
            return nearest
        else:
            return nearest[1]
 
    def find_nearest_dist_nodes(self, point, distance):
        """!
        @brief Find neighbors that are located in area that is covered by specified distance.
 
        @param[in] point (list): Coordinates that is considered as centroid for searching.
        @param[in] distance (double): Distance from the center where searching is performed.
 
        @return (list) Neighbors in area that is specified by point (center) and distance (radius).
 
        """
 
        best_nodes = []
        if self._root is not None:
            self.__recursive_nearest_nodes(point, distance, distance * distance, self._root, best_nodes)
 
        return best_nodes
 
    def __recursive_nearest_nodes(self, point, distance, sqrt_distance, node_head, best_nodes):
        """!
        @brief Returns list of neighbors such as tuple (distance, node) that is located in area that is covered by distance.
 
        @param[in] point (list): Coordinates that is considered as centroid for searching
        @param[in] distance (double): Distance from the center where searching is performed.
        @param[in] sqrt_distance (double): Square distance from the center where searching is performed.
        @param[in] node_head (node): Node from that searching is performed.
        @param[in|out] best_nodes (list): List of founded nodes.
 
        """
 
        if node_head.right is not None:
            minimum = node_head.data[node_head.disc] - distance
            if point[node_head.disc] >= minimum:
                self.__recursive_nearest_nodes(point, distance, sqrt_distance, node_head.right, best_nodes)
 
        if node_head.left is not None:
            maximum = node_head.data[node_head.disc] + distance
            if point[node_head.disc] < maximum:
                self.__recursive_nearest_nodes(point, distance, sqrt_distance, node_head.left, best_nodes)
 
        candidate_distance = euclidean_distance_square(point, node_head.data)
        if candidate_distance <= sqrt_distance:
            best_nodes.append((candidate_distance, node_head))
 
 
 
class kdtree(kdtree_balanced):
    """!
    @brief   Represents KD Tree that is a space-partitioning data structure for organizing points
              in a k-dimensional space.
    @details In the term k-d tree, k denotes the dimensionality of the space being represented. Each data point is
              represented as a node in the k-d tree in the form of a record of type node. The tree supports
              dynamic construction when nodes can be dynamically added and removed. As a result KD-tree might not be
              balanced if methods `insert` and `remove` are used to built the tree. If the tree is built using
              constructor where all points are passed to the tree then balanced tree is built. Single point search and
              range-search procedures have complexity is `O(n)` in worse case in case of unbalanced tree.
              If there is no need to build dynamic KD-tree, then it is much better to use static KD-tree
              `kdtree_balanced`.
    
    There is an example how to use KD-tree to search nodes (points from input data) that are nearest to some point:
    @code
        # Import required modules
        from pyclustering.samples.definitions import SIMPLE_SAMPLES;
        from pyclustering.container.kdtree import kdtree;
        from pyclustering.utils import read_sample;
        
        # Read data from text file
        sample = read_sample(SIMPLE_SAMPLES.SAMPLE_SIMPLE3);
        
        # Create instance of KD-tree and initialize (fill) it by read data.
        # In this case the tree is balanced.
        tree_instance = kdtree(sample);
        
        # Search for nearest point
        search_distance = 0.3;
        nearest_node = tree_instance.find_nearest_dist_node([1.12, 4.31], search_distance);
        
        # Search for nearest point in radius 0.3
        nearest_nodes = tree_instance.find_nearest_dist_nodes([1.12, 4.31], search_distance);
        print("Nearest nodes:", nearest_nodes);
    @endcode
 
    In case of building KD-tree using `insert` and `remove` method, the output KD-tree might be unbalanced - here
    is an example that demonstrates this:
    @code
        from pyclustering.container.kdtree import kdtree, kdtree_visualizer
        from pyclustering.utils import read_sample
        from pyclustering.samples.definitions import FCPS_SAMPLES
 
        sample = read_sample(FCPS_SAMPLES.SAMPLE_TWO_DIAMONDS)
 
        # Build tree using constructor - balanced will be built because tree will know about all points.
        tree_instance = kdtree(sample)
        kdtree_visualizer(tree_instance).visualize()
 
        # Build tree using `insert` only - unbalanced tree will be built.
        tree_instance = kdtree()
        for point in sample:
            tree_instance.insert(point)
 
        kdtree_visualizer(tree_instance).visualize()
    @endcode
 
    There are two figures where difference between balanced and unbalanced KD-trees is demonstrated.
 
    @image html kd_tree_unbalanced_two_diamonds.png "Fig. 1. Balanced KD-tree for sample 'TwoDiamonds'."
    @image html kd_tree_unbalanced_two_diamonds.png "Fig. 2. Unbalanced KD-tree for sample 'TwoDiamonds'."
 
    @see kdtree_balanced
 
    """
    
    def __init__(self, data_list=None, payload_list=None):
        """!
        @brief Create kd-tree from list of points and from according list of payloads.
        @details If lists were not specified then empty kd-tree will be created.
        
        @param[in] data_list (list): Insert points from the list to created KD tree.
        @param[in] payload_list (list): Insert payload from the list to created KD tree, length should be equal to
                    length of data_list if it is specified.
        
        """
 
        super().__init__(data_list, payload_list)
 
    def insert(self, point, payload=None):
        """!
        @brief Insert new point with payload to kd-tree.
        
        @param[in] point (list): Coordinates of the point of inserted node.
        @param[in] payload (any-type): Payload of inserted node. It can be ID of the node or
                    some useful payload that belongs to the point.
        
        @return (node) Inserted node to the kd-tree.
        
        """
        
        if self._root is None:
            self._dimension = len(point)
            self._root = node(point, payload, None, None, 0)
            self._point_comparator = self._create_point_comparator(type(point))
 
            self._length += 1
            return self._root
 
        cur_node = self._root
 
        while True:
            discriminator = (cur_node.disc + 1) % self._dimension
 
            if cur_node.data[cur_node.disc] <= point[cur_node.disc]:
                if cur_node.right is None:
                    cur_node.right = node(point, payload, None, None, discriminator, cur_node)
 
                    self._length += 1
                    return cur_node.right
                else: 
                    cur_node = cur_node.right
            
            else:
                if cur_node.left is None:
                    cur_node.left = node(point, payload, None, None, discriminator, cur_node)
 
                    self._length += 1
                    return cur_node.left
                else:
                    cur_node = cur_node.left
 
    def remove(self, point, **kwargs):
        """!
        @brief Remove specified point from kd-tree.
        @details It removes the first found node that satisfy to the input parameters. Make sure that
                  pair (point, payload) is unique for each node, otherwise the first found is removed.
        
        @param[in] point (list): Coordinates of the point of removed node.
        @param[in] **kwargs: Arbitrary keyword arguments (available arguments: 'payload').
        
        <b>Keyword Args:</b><br>
            - payload (any): Payload of the node that should be removed.
        
        @return (node) Root if node has been successfully removed, otherwise None.
        
        """
 
        if 'payload' in kwargs:
            node_for_remove = self.find_node_with_payload(point, kwargs['payload'], None)
        else:
            node_for_remove = self.find_node(point, None)
        
        if node_for_remove is None:
            return None
 
        self._length -= 1
 
        parent = node_for_remove.parent
        minimal_node = self.__recursive_remove(node_for_remove)
        if parent is None:
            self._root = minimal_node
            
            # If all k-d tree was destroyed
            if minimal_node is not None:
                minimal_node.parent = None
        else:
            if parent.left is node_for_remove:
                parent.left = minimal_node
            elif parent.right is node_for_remove:
                parent.right = minimal_node
        
        return self._root
 
    def __recursive_remove(self, node_removed):
        """!
        @brief Delete node and return root of subtree.
        
        @param[in] node_removed (node): Node that should be removed.
        
        @return (node) Minimal node in line with coordinate that is defined by discriminator.
        
        """
                
        # Check if it is leaf
        if (node_removed.right is None) and (node_removed.left is None):
            return None
        
        discriminator = node_removed.disc
        
        # Check if only left branch exist
        if node_removed.right is None:
            node_removed.right = node_removed.left
            node_removed.left = None
        
        # Find minimal node in line with coordinate that is defined by discriminator
        minimal_node = self.__find_minimal_node(node_removed.right, discriminator)
        parent = minimal_node.parent
        
        if parent.left is minimal_node:
            parent.left = self.__recursive_remove(minimal_node)
        elif parent.right is minimal_node:
            parent.right = self.__recursive_remove(minimal_node)
        
        minimal_node.parent = node_removed.parent
        minimal_node.disc = node_removed.disc
        minimal_node.right = node_removed.right
        minimal_node.left = node_removed.left
        
        # Update parent for successors of previous parent.
        if minimal_node.right is not None:
            minimal_node.right.parent = minimal_node
 
        if minimal_node.left is not None:
            minimal_node.left.parent = minimal_node
 
        return minimal_node
 
    def __find_minimal_node(self, node_head, discriminator):
        """!
        @brief Find minimal node in line with coordinate that is defined by discriminator.
        
        @param[in] node_head (node): Node of KD tree from that search should be started.
        @param[in] discriminator (uint): Coordinate number that is used for comparison.
        
        @return (node) Minimal node in line with discriminator from the specified node.
        
        """
        
        min_key = lambda cur_node: cur_node.data[discriminator]
 
        stack, candidates = [], []
        is_finished = False
 
        while is_finished is False:
            if node_head is not None:
                stack.append(node_head)
                node_head = node_head.left
            else:
                if len(stack) != 0:
                    node_head = stack.pop()
                    candidates.append(node_head)
                    node_head = node_head.right
                else:
                    is_finished = True
 
        return min(candidates, key=min_key)