xref: /external/chromium_org/third_party/bintrees/bintrees/treemixin.py

#!/usr/bin/env python
#coding:utf-8
# Author:  Mozman
# Purpose: treemixin provides top level functions for binary trees
# Created: 03.05.2010
# Copyright (c) 2010-2013 by Manfred Moitzi
# License: MIT License

from __future__ import absolute_import

from .walker import Walker
from .treeslice import TreeSlice

class TreeMixin(object):
    """
    Abstract-Base-Class for the pure Python Trees: BinaryTree, AVLTree and RBTree
    Mixin-Class for the Cython based Trees: FastBinaryTree, FastAVLTree, FastRBTree

    The TreeMixin Class
    ===================

    T has to implement following properties
    ---------------------------------------

    count -- get node count

    T has to implement following methods
    ------------------------------------

    get_walker(...)
        get a tree walker object, provides methods to traverse the tree see walker.py

    insert(...)
        insert(key, value) <==> T[key] = value, insert key into T

    remove(...)
        remove(key) <==> del T[key], remove key from T

    clear(...)
        T.clear() -> None.  Remove all items from T.

    Methods defined here
    --------------------

    * __contains__(k) -> True if T has a key k, else False, O(log(n))
    * __delitem__(y) <==> del T[y], del T[s:e], O(log(n))
    * __getitem__(y) <==> T[y], T[s:e], O(log(n))
    * __iter__() <==> iter(T)
    * __len__() <==> len(T), O(1)
    * __max__() <==> max(T), get max item (k,v) of T, O(log(n))
    * __min__() <==> min(T), get min item (k,v) of T, O(log(n))
    * __and__(other) <==> T & other, intersection
    * __or__(other) <==> T | other, union
    * __sub__(other) <==> T - other, difference
    * __xor__(other) <==> T ^ other, symmetric_difference
    * __repr__() <==> repr(T)
    * __setitem__(k, v) <==> T[k] = v, O(log(n))
    * clear() -> None, Remove all items from T, , O(n)
    * copy() -> a shallow copy of T, O(n*log(n))
    * discard(k) -> None, remove k from T, if k is present, O(log(n))
    * get(k[,d]) -> T[k] if k in T, else d, O(log(n))
    * is_empty() -> True if len(T) == 0, O(1)
    * items([reverse]) -> generator for (k, v) items of T, O(n)
    * keys([reverse]) -> generator for keys of T, O(n)
    * values([reverse]) -> generator for values of  T, O(n)
    * pop(k[,d]) -> v, remove specified key and return the corresponding value, O(log(n))
    * popitem() -> (k, v), remove and return some (key, value) pair as a 2-tuple, O(log(n))
    * setdefault(k[,d]) -> T.get(k, d), also set T[k]=d if k not in T, O(log(n))
    * update(E) -> None.  Update T from dict/iterable E, O(E*log(n))
    * foreach(f, [order]) -> visit all nodes of tree and call f(k, v) for each node, O(n)

    slicing by keys

    * itemslice(s, e) -> generator for (k, v) items of T for s <= key < e, O(n)
    * keyslice(s, e) -> generator for keys of T for s <= key < e, O(n)
    * valueslice(s, e) -> generator for values of T for s <= key < e, O(n)
    * T[s:e] -> TreeSlice object, with keys in range s <= key < e, O(n)
    * del T[s:e] -> remove items by key slicing, for s <= key < e, O(n)

    if 's' is None or T[:e] TreeSlice/iterator starts with value of min_key()
    if 'e' is None or T[s:] TreeSlice/iterator ends with value of max_key()
    T[:] is a TreeSlice which represents the whole tree.

    TreeSlice is a tree wrapper with range check, and contains no references
    to objects, deleting objects in the associated tree also deletes the object
    in the TreeSlice.

    * TreeSlice[k] -> get value for key k, raises KeyError if k not exists in range s:e
    * TreeSlice[s1:e1] -> TreeSlice object, with keys in range s1 <= key < e1

      * new lower bound is max(s, s1)
      * new upper bound is min(e, e1)

    TreeSlice methods:

    * items() -> generator for (k, v) items of T, O(n)
    * keys() -> generator for keys of T, O(n)
    * values() -> generator for values of  T, O(n)
    * __iter__ <==> keys()
    * __repr__ <==> repr(T)
    * __contains__(key)-> True if TreeSlice has a key k, else False, O(log(n))

    prev/succ operations

    * prev_item(key) -> get (k, v) pair, where k is predecessor to key, O(log(n))
    * prev_key(key) -> k, get the predecessor of key, O(log(n))
    * succ_item(key) -> get (k,v) pair as a 2-tuple, where k is successor to key, O(log(n))
    * succ_key(key) -> k, get the successor of key, O(log(n))
    * floor_item(key) -> get (k, v) pair, where k is the greatest key less than or equal to key, O(log(n))
    * floor_key(key) -> k, get the greatest key less than or equal to key, O(log(n))
    * ceiling_item(key) -> get (k, v) pair, where k is the smallest key greater than or equal to key, O(log(n))
    * ceiling_key(key) -> k, get the smallest key greater than or equal to key, O(log(n))

    Heap methods

    * max_item() -> get largest (key, value) pair of T, O(log(n))
    * max_key() -> get largest key of T, O(log(n))
    * min_item() -> get smallest (key, value) pair of T, O(log(n))
    * min_key() -> get smallest key of T, O(log(n))
    * pop_min() -> (k, v), remove item with minimum key, O(log(n))
    * pop_max() -> (k, v), remove item with maximum key, O(log(n))
    * nlargest(i[,pop]) -> get list of i largest items (k, v), O(i*log(n))
    * nsmallest(i[,pop]) -> get list of i smallest items (k, v), O(i*log(n))

    Set methods (using frozenset)

    * intersection(t1, t2, ...) -> Tree with keys *common* to all trees
    * union(t1, t2, ...) -> Tree with keys from *either* trees
    * difference(t1, t2, ...) -> Tree with keys in T but not any of t1, t2, ...
    * symmetric_difference(t1) -> Tree with keys in either T and t1  but not both
    * issubset(S) -> True if every element in T is in S
    * issuperset(S) -> True if every element in S is in T
    * isdisjoint(S) ->  True if T has a null intersection with S

    Classmethods

    * fromkeys(S[,v]) -> New tree with keys from S and values equal to v.

    """
    def get_walker(self):
        return Walker(self)

    def __repr__(self):
        """ x.__repr__(...) <==> repr(x) """
        tpl = "%s({%s})" % (self.__class__.__name__, '%s')
        return tpl % ", ".join( ("%r: %r" % item for item in self.items()) )

    def copy(self):
        """ T.copy() -> get a shallow copy of T. """
        tree = self.__class__()
        self.foreach(tree.insert, order=-1)
        return tree
    __copy__ = copy

    def __contains__(self, key):
        """ k in T -> True if T has a key k, else False """
        try:
            self.get_value(key)
            return True
        except KeyError:
            return False

    def __len__(self):
        """ x.__len__() <==> len(x) """
        return self.count

    def __min__(self):
        """ x.__min__() <==> min(x) """
        return self.min_item()

    def __max__(self):
        """ x.__max__() <==> max(x) """
        return self.max_item()

    def __and__(self, other):
        """ x.__and__(other) <==> self & other """
        return self.intersection(other)

    def __or__(self, other):
        """ x.__or__(other) <==> self | other """
        return self.union(other)

    def __sub__(self, other):
        """ x.__sub__(other) <==> self - other """
        return self.difference(other)

    def __xor__(self, other):
        """ x.__xor__(other) <==> self ^ other """
        return self.symmetric_difference(other)

    def discard(self, key):
        """ x.discard(k) -> None, remove k from T, if k is present """
        try:
            self.remove(key)
        except KeyError:
            pass

    def __del__(self):
        self.clear()

    def is_empty(self):
        """ x.is_empty() -> False if T contains any items else True"""
        return self.count == 0

    def keys(self, reverse=False):
        """ T.iterkeys([reverse]) -> an iterator over the keys of T, in ascending
        order if reverse is True, iterate in descending order, reverse defaults
        to False
        """
        return ( item[0] for item in self.items(reverse) )
    __iter__ = keys

    def __reversed__(self):
        return self.keys(reverse=True)

    def values(self, reverse=False):
        """ T.values([reverse]) -> an iterator over the values of T, in ascending order
        if reverse is True, iterate in descending order, reverse defaults to False
        """
        return ( item[1] for item in self.items(reverse) )

    def items(self, reverse=False):
        """ T.items([reverse]) -> an iterator over the (key, value) items of T,
        in ascending order if reverse is True, iterate in descending order,
        reverse defaults to False
        """
        if self.is_empty():
            return []

        def default_iterator(node):
            direction = 1 if reverse else 0
            other = 1 - direction
            go_down = True
            while True:
                if node.has_child(direction) and go_down:
                    node.push()
                    node.down(direction)
                else:
                    yield node.item
                    if node.has_child(other):
                        node.down(other)
                        go_down = True
                    else:
                        if node.stack_is_empty():
                            return  # all done
                        node.pop()
                        go_down = False

        treewalker = self.get_walker()
        try:  # specialized iterators
            if reverse:
                return treewalker.iter_items_backward()
            else:
                return treewalker.iter_items_forward()
        except AttributeError:
            return default_iterator(treewalker)

    def __getitem__(self, key):
        """ x.__getitem__(y) <==> x[y] """
        if isinstance(key, slice):
            return TreeSlice(self, key.start, key.stop)
        else:
            return self.get_value(key)

    def __setitem__(self, key, value):
        """ x.__setitem__(i, y) <==> x[i]=y """
        if isinstance(key, slice):
            raise ValueError('setslice is not supported')
        self.insert(key, value)

    def __delitem__(self, key):
        """ x.__delitem__(y) <==> del x[y] """
        if isinstance(key, slice):
            self.delitems(self.keyslice(key.start, key.stop))
        else:
            self.remove(key)

    def delitems(self, keys):
        """ T.delitems(keys) -> remove all items by keys
        """
        # convert generator to a set, because the content of the
        # tree will be modified!
        for key in frozenset(keys):
            self.remove(key)

    def keyslice(self, startkey, endkey):
        """ T.keyslice(startkey, endkey) -> key iterator:
        startkey <= key < endkey.
        """
        return ( item[0] for item in self.itemslice(startkey, endkey) )

    def itemslice(self, startkey, endkey):
        """ T.itemslice(s, e) -> item iterator: s <= key < e.

        if s is None: start with min element -> T[:e]
        if e is None: end with max element -> T[s:]
        T[:] -> all items

        """
        if self.is_empty():
            return

        if startkey is None:
            # no lower bound
            can_go_left = lambda: node.has_left() and visit_left
        else:
            # don't visit subtrees without keys in search range
            can_go_left = lambda: node.key > startkey and node.has_left() and visit_left

        if endkey is None:
            # no upper bound
            can_go_right = lambda: node.has_right()
        else:
            # don't visit subtrees without keys in search range
            can_go_right = lambda: node.key < endkey and node.has_right()

        if (startkey, endkey) == (None, None):
            key_in_range = lambda: True
        elif startkey is None:
            key_in_range = lambda: node.key < endkey
        elif endkey is None:
            key_in_range = lambda: node.key >= startkey
        else:
            key_in_range = lambda: startkey <= node.key < endkey

        node = self.get_walker()
        visit_left = True
        while True:
            if can_go_left():
                node.push()
                node.go_left()
            else:
                if key_in_range():
                    yield node.item
                if can_go_right():
                    node.go_right()
                    visit_left = True
                else:
                    if node.stack_is_empty():
                        return
                    node.pop()
                    # left side is already done
                    visit_left = False

    def valueslice(self, startkey, endkey):
        """ T.valueslice(startkey, endkey) -> value iterator:
        startkey <= key < endkey.
        """
        return ( item[1] for item in self.itemslice(startkey, endkey) )

    def get_value(self, key):
        node = self.root
        while node is not None:
            if key == node.key:
                return node.value
            elif key < node.key:
                node = node.left
            else:
                node = node.right
        raise KeyError(str(key))

    def __getstate__(self):
        return dict(self.items())

    def __setstate__(self, state):
        # note for myself: this is called like __init__, so don't use clear()
        # to remove existing data!
        self._root = None
        self._count = 0
        self.update(state)

    def setdefault(self, key, default=None):
        """ T.setdefault(k[,d]) -> T.get(k,d), also set T[k]=d if k not in T """
        try:
            return self.get_value(key)
        except KeyError:
            self.insert(key, default)
            return default

    def update(self, *args):
        """ T.update(E) -> None. Update T from E : for (k, v) in E: T[k] = v """
        for items in args:
            try:
                generator = items.items()
            except AttributeError:
                generator = iter(items)

            for key, value in generator:
                self.insert(key, value)

    @classmethod
    def fromkeys(cls, iterable, value=None):
        """ T.fromkeys(S[,v]) -> New tree with keys from S and values equal to v.

        v defaults to None.
        """
        tree = cls()
        for key in iterable:
            tree.insert(key, value)
        return tree

    def get(self, key, default=None):
        """ T.get(k[,d]) -> T[k] if k in T, else d.  d defaults to None. """
        try:
            return self.get_value(key)
        except KeyError:
            return default

    def pop(self, key, *args):
        """ T.pop(k[,d]) -> v, remove specified key and return the corresponding value
        If key is not found, d is returned if given, otherwise KeyError is raised
        """
        if len(args) > 1:
            raise TypeError("pop expected at most 2 arguments, got %d" % (1 + len(args)))
        try:
            value = self.get_value(key)
            self.remove(key)
            return value
        except KeyError:
            if len(args) == 0:
                raise
            else:
                return args[0]

    def popitem(self):
        """ T.popitem() -> (k, v), remove and return some (key, value) pair as a
        2-tuple; but raise KeyError if T is empty
        """
        if self.is_empty():
            raise KeyError("popitem(): tree is empty")
        walker = self.get_walker()
        walker.goto_leaf()
        result = walker.item
        self.remove(walker.key)
        return result

    def foreach(self, func, order=0):
        """ Visit all tree nodes and process key, value.

        parm func: function(key, value)
        param int order: inorder = 0, preorder = -1, postorder = +1

        """
        def _traverse():
            if order == -1:
                func(node.key, node.value)
            if node.has_left():
                node.push()
                node.go_left()
                _traverse()
                node.pop()
            if order == 0:
                func(node.key, node.value)
            if node.has_right():
                node.push()
                node.go_right()
                _traverse()
                node.pop()
            if order == +1:
                func(node.key, node.value)

        node = self.get_walker()
        _traverse()

    def min_item(self):
        """ Get item with min key of tree, raises ValueError if tree is empty. """
        if self.count == 0:
            raise ValueError("Tree is empty")
        node = self._root
        while node.left is not None:
            node = node.left
        return (node.key, node.value)

    def pop_min(self):
        """ T.pop_min() -> (k, v), remove item with minimum key, raise ValueError
        if T is empty.
        """
        item = self.min_item()
        self.remove(item[0])
        return item

    def min_key(self):
        """ Get min key of tree, raises ValueError if tree is empty. """
        key, value = self.min_item()
        return key

    def prev_item(self, key):
        """ Get predecessor (k,v) pair of key, raises KeyError if key is min key
        or key does not exist.
        """
        if self.count == 0:
            raise KeyError("Tree is empty")
        walker = self.get_walker()
        return walker.prev_item(key)

    def succ_item(self, key):
        """ Get successor (k,v) pair of key, raises KeyError if key is max key
        or key does not exist.
        """
        if self.count == 0:
            raise KeyError("Tree is empty")
        walker = self.get_walker()
        return walker.succ_item(key)

    def prev_key(self, key):
        """ Get predecessor to key, raises KeyError if key is min key
        or key does not exist.
        """
        key, value = self.prev_item(key)
        return key

    def succ_key(self, key):
        """ Get successor to key, raises KeyError if key is max key
        or key does not exist.
        """
        key, value = self.succ_item(key)
        return key

    def floor_item(self, key):
        """ Get the element (k,v) pair associated with the greatest key less
        than or equal to the given key, raises KeyError if there is no such key.
        """
        if self.count == 0:
            raise KeyError("Tree is empty")
        walker = self.get_walker()
        return walker.floor_item(key)

    def floor_key(self, key):
        """ Get the greatest key less than or equal to the given key, raises
        KeyError if there is no such key.
        """
        key, value = self.floor_item(key)
        return key

    def ceiling_item(self, key):
        """ Get the element (k,v) pair associated with the smallest key greater
        than or equal to the given key, raises KeyError if there is no such key.
        """
        if self.count == 0:
            raise KeyError("Tree is empty")
        walker = self.get_walker()
        return walker.ceiling_item(key)

    def ceiling_key(self, key):
        """ Get the smallest key greater than or equal to the given key, raises
        KeyError if there is no such key.
        """
        key, value = self.ceiling_item(key)
        return key

    def max_item(self):
        """ Get item with max key of tree, raises ValueError if tree is empty. """
        if self.count == 0:
            raise ValueError("Tree is empty")
        node = self._root
        while node.right is not None:
            node = node.right
        return (node.key, node.value)

    def pop_max(self):
        """ T.pop_max() -> (k, v), remove item with maximum key, raise ValueError
        if T is empty.
        """
        item = self.max_item()
        self.remove(item[0])
        return item

    def max_key(self):
        """ Get max key of tree, raises ValueError if tree is empty. """
        key, value = self.max_item()
        return key

    def nsmallest(self, n, pop=False):
        """ T.nsmallest(n) -> get list of n smallest items (k, v).
        If pop is True, remove items from T.
        """
        if pop:
            return [self.pop_min() for _ in range(min(len(self), n))]
        else:
            items = self.items()
            return  [ next(items) for _ in range(min(len(self), n)) ]

    def nlargest(self, n, pop=False):
        """ T.nlargest(n) -> get list of n largest items (k, v).
        If pop is True remove items from T.
        """
        if pop:
            return [self.pop_max() for _ in range(min(len(self), n))]
        else:
            items = self.items(reverse=True)
            return [ next(items) for _ in range(min(len(self), n)) ]

    def intersection(self, *trees):
        """ x.intersection(t1, t2, ...) -> Tree, with keys *common* to all trees
        """
        thiskeys = frozenset(self.keys())
        sets = _build_sets(trees)
        rkeys = thiskeys.intersection(*sets)
        return self.__class__( ((key, self.get(key)) for key in rkeys) )

    def union(self, *trees):
        """ x.union(t1, t2, ...) -> Tree with keys from *either* trees
        """
        thiskeys = frozenset(self.keys())
        rkeys = thiskeys.union(*_build_sets(trees))
        return self.__class__( ((key, self.get(key)) for key in rkeys) )

    def difference(self, *trees):
        """ x.difference(t1, t2, ...) -> Tree with keys in T but not any of t1,
        t2, ...
        """
        thiskeys = frozenset(self.keys())
        rkeys = thiskeys.difference(*_build_sets(trees))
        return self.__class__( ((key, self.get(key)) for key in rkeys) )

    def symmetric_difference(self, tree):
        """ x.symmetric_difference(t1) -> Tree with keys in either T and t1  but
        not both
        """
        thiskeys = frozenset(self.keys())
        rkeys = thiskeys.symmetric_difference(frozenset(tree.keys()))
        return self.__class__( ((key, self.get(key)) for key in rkeys) )

    def issubset(self, tree):
        """ x.issubset(tree) -> True if every element in x is in tree """
        thiskeys = frozenset(self.keys())
        return thiskeys.issubset(frozenset(tree.keys()))

    def issuperset(self, tree):
        """ x.issubset(tree) -> True if every element in tree is in x """
        thiskeys = frozenset(self.keys())
        return thiskeys.issuperset(frozenset(tree.keys()))

    def isdisjoint(self, tree):
        """ x.isdisjoint(S) ->  True if x has a null intersection with tree """
        thiskeys = frozenset(self.keys())
        return thiskeys.isdisjoint(frozenset(tree.keys()))


def _build_sets(trees):
    return [ frozenset(tree.keys()) for tree in trees ]