Source code for astropy.table.bst

# Licensed under a 3-clause BSD style license - see LICENSE.rst
import operator

__all__ = ["BST"]


class MaxValue:
    """
    Represents an infinite value for purposes
    of tuple comparison.
    """

    def __gt__(self, other):
        return True

    def __ge__(self, other):
        return True

    def __lt__(self, other):
        return False

    def __le__(self, other):
        return False

    def __repr__(self):
        return "MAX"

    __str__ = __repr__


class MinValue:
    """
    The opposite of MaxValue, i.e. a representation of
    negative infinity.
    """

    def __lt__(self, other):
        return True

    def __le__(self, other):
        return True

    def __gt__(self, other):
        return False

    def __ge__(self, other):
        return False

    def __repr__(self):
        return "MIN"

    __str__ = __repr__


class Epsilon:
    """
    Represents the "next largest" version of a given value,
    so that for all valid comparisons we have
    x < y < Epsilon(y) < z whenever x < y < z and x, z are
    not Epsilon objects.

    Parameters
    ----------
    val : object
        Original value
    """

    __slots__ = ("val",)

    def __init__(self, val):
        self.val = val

    def __lt__(self, other):
        if self.val == other:
            return False
        return self.val < other

    def __gt__(self, other):
        if self.val == other:
            return True
        return self.val > other

    def __eq__(self, other):
        return False

    def __repr__(self):
        return repr(self.val) + " + epsilon"


class Node:
    """
    An element in a binary search tree, containing
    a key, data, and references to children nodes and
    a parent node.

    Parameters
    ----------
    key : tuple
        Node key
    data : list or int
        Node data
    """

    __lt__ = lambda x, y: x.key < y.key
    __le__ = lambda x, y: x.key <= y.key
    __eq__ = lambda x, y: x.key == y.key
    __ge__ = lambda x, y: x.key >= y.key
    __gt__ = lambda x, y: x.key > y.key
    __ne__ = lambda x, y: x.key != y.key
    __slots__ = ("key", "data", "left", "right")

    # each node has a key and data list
    def __init__(self, key, data):
        self.key = key
        self.data = data if isinstance(data, list) else [data]
        self.left = None
        self.right = None

    def replace(self, child, new_child):
        """
        Replace this node's child with a new child.
        """
        if self.left is not None and self.left == child:
            self.left = new_child
        elif self.right is not None and self.right == child:
            self.right = new_child
        else:
            raise ValueError("Cannot call replace() on non-child")

    def remove(self, child):
        """
        Remove the given child.
        """
        self.replace(child, None)

    def set(self, other):
        """
        Copy the given node.
        """
        self.key = other.key
        self.data = other.data[:]

    def __str__(self):
        return str((self.key, self.data))

    def __repr__(self):
        return str(self)


[docs] class BST: """ A basic binary search tree in pure Python, used as an engine for indexing. Parameters ---------- data : Table Sorted columns of the original table row_index : Column object Row numbers corresponding to data columns unique : bool Whether the values of the index must be unique. Defaults to False. """ NodeClass = Node def __init__(self, data, row_index, unique=False): self.root = None self.size = 0 self.unique = unique for key, row in zip(data, row_index): self.add(tuple(key), row)
[docs] def add(self, key, data=None): """ Add a key, data pair. """ if data is None: data = key self.size += 1 node = self.NodeClass(key, data) curr_node = self.root if curr_node is None: self.root = node return while True: if node < curr_node: if curr_node.left is None: curr_node.left = node break curr_node = curr_node.left elif node > curr_node: if curr_node.right is None: curr_node.right = node break curr_node = curr_node.right elif self.unique: raise ValueError("Cannot insert non-unique value") else: # add data to node curr_node.data.extend(node.data) curr_node.data = sorted(curr_node.data) return
[docs] def find(self, key): """ Return all data values corresponding to a given key. Parameters ---------- key : tuple Input key Returns ------- data_vals : list List of rows corresponding to the input key """ node, parent = self.find_node(key) return node.data if node is not None else []
[docs] def find_node(self, key): """ Find the node associated with the given key. """ if self.root is None: return (None, None) return self._find_recursive(key, self.root, None)
[docs] def shift_left(self, row): """ Decrement all rows larger than the given row. """ for node in self.traverse(): node.data = [x - 1 if x > row else x for x in node.data]
[docs] def shift_right(self, row): """ Increment all rows greater than or equal to the given row. """ for node in self.traverse(): node.data = [x + 1 if x >= row else x for x in node.data]
def _find_recursive(self, key, node, parent): try: if key == node.key: return (node, parent) elif key > node.key: if node.right is None: return (None, None) return self._find_recursive(key, node.right, node) else: if node.left is None: return (None, None) return self._find_recursive(key, node.left, node) except TypeError: # wrong key type return (None, None)
[docs] def traverse(self, order="inorder"): """ Return nodes of the BST in the given order. Parameters ---------- order : str The order in which to recursively search the BST. Possible values are: "preorder": current node, left subtree, right subtree "inorder": left subtree, current node, right subtree "postorder": left subtree, right subtree, current node """ if order == "preorder": return self._preorder(self.root, []) elif order == "inorder": return self._inorder(self.root, []) elif order == "postorder": return self._postorder(self.root, []) raise ValueError(f'Invalid traversal method: "{order}"')
[docs] def items(self): """ Return BST items in order as (key, data) pairs. """ return [(x.key, x.data) for x in self.traverse()]
[docs] def sort(self): """ Make row order align with key order. """ i = 0 for node in self.traverse(): num_rows = len(node.data) node.data = list(range(i, i + num_rows)) i += num_rows
[docs] def sorted_data(self): """ Return BST rows sorted by key values. """ return [x for node in self.traverse() for x in node.data]
def _preorder(self, node, lst): if node is None: return lst lst.append(node) self._preorder(node.left, lst) self._preorder(node.right, lst) return lst def _inorder(self, node, lst): if node is None: return lst self._inorder(node.left, lst) lst.append(node) self._inorder(node.right, lst) return lst def _postorder(self, node, lst): if node is None: return lst self._postorder(node.left, lst) self._postorder(node.right, lst) lst.append(node) return lst def _substitute(self, node, parent, new_node): if node is self.root: self.root = new_node else: parent.replace(node, new_node)
[docs] def remove(self, key, data=None): """ Remove data corresponding to the given key. Parameters ---------- key : tuple The key to remove data : int or None If None, remove the node corresponding to the given key. If not None, remove only the given data value from the node. Returns ------- successful : bool True if removal was successful, false otherwise """ node, parent = self.find_node(key) if node is None: return False if data is not None: if data not in node.data: raise ValueError("Data does not belong to correct node") elif len(node.data) > 1: node.data.remove(data) return True if node.left is None and node.right is None: self._substitute(node, parent, None) elif node.left is None and node.right is not None: self._substitute(node, parent, node.right) elif node.right is None and node.left is not None: self._substitute(node, parent, node.left) else: # find largest element of left subtree curr_node = node.left parent = node while curr_node.right is not None: parent = curr_node curr_node = curr_node.right self._substitute(curr_node, parent, curr_node.left) node.set(curr_node) self.size -= 1 return True
[docs] def is_valid(self): """ Returns whether this is a valid BST. """ return self._is_valid(self.root)
def _is_valid(self, node): if node is None: return True return ( (node.left is None or node.left <= node) and (node.right is None or node.right >= node) and self._is_valid(node.left) and self._is_valid(node.right) )
[docs] def range(self, lower, upper, bounds=(True, True)): """ Return all nodes with keys in the given range. Parameters ---------- lower : tuple Lower bound upper : tuple Upper bound bounds : (2,) tuple of bool Indicates whether the search should be inclusive or exclusive with respect to the endpoints. The first argument corresponds to an inclusive lower bound, and the second argument to an inclusive upper bound. """ nodes = self.range_nodes(lower, upper, bounds) return [x for node in nodes for x in node.data]
[docs] def range_nodes(self, lower, upper, bounds=(True, True)): """ Return nodes in the given range. """ if self.root is None: return [] # op1 is <= or <, op2 is >= or > op1 = operator.le if bounds[0] else operator.lt op2 = operator.ge if bounds[1] else operator.gt return self._range(lower, upper, op1, op2, self.root, [])
[docs] def same_prefix(self, val): """ Assuming the given value has smaller length than keys, return nodes whose keys have this value as a prefix. """ if self.root is None: return [] nodes = self._same_prefix(val, self.root, []) return [x for node in nodes for x in node.data]
def _range(self, lower, upper, op1, op2, node, lst): if op1(lower, node.key) and op2(upper, node.key): lst.append(node) if upper > node.key and node.right is not None: self._range(lower, upper, op1, op2, node.right, lst) if lower < node.key and node.left is not None: self._range(lower, upper, op1, op2, node.left, lst) return lst def _same_prefix(self, val, node, lst): prefix = node.key[: len(val)] if prefix == val: lst.append(node) if prefix <= val and node.right is not None: self._same_prefix(val, node.right, lst) if prefix >= val and node.left is not None: self._same_prefix(val, node.left, lst) return lst def __repr__(self): return f"<{self.__class__.__name__}>" def _print(self, node, level): line = "\t" * level + str(node) + "\n" if node.left is not None: line += self._print(node.left, level + 1) if node.right is not None: line += self._print(node.right, level + 1) return line @property def height(self): """ Return the BST height. """ return self._height(self.root) def _height(self, node): if node is None: return -1 return max(self._height(node.left), self._height(node.right)) + 1
[docs] def replace_rows(self, row_map): """ Replace all rows with the values they map to in the given dictionary. Any rows not present as keys in the dictionary will have their nodes deleted. Parameters ---------- row_map : dict Mapping of row numbers to new row numbers """ for _, data in self.items(): # noqa: PERF102 data[:] = [row_map[x] for x in data if x in row_map]