arvoreb.py

#!/usr/bin/python
# =============================================================================
#                                                 Pesquisa e Ordenacao - 2016.1
# =============================================================================
import random
import bisect

class _BTreeNode(object):
    def __init__(self, values=None, children=None):
        self.parent = None
        self.values = values or []
        self.children = children
        if self.children:
            for i in self.children:
                i.parent = self

    def __str__(self):
        return 'No\' %r (%d filhos)' % ( self.values , len(self.children) if self.children else 0 )
    #           id(self), id(self.parent),

    def pretty_print(self, tab=''):
        print('%s%s' % (tab, self))
        if self.children:
            for i in self.children:
                i.pretty_print(tab + '   ')

    def check_valid(self, tree): # checa integridade da arvore
        innerNode = self.children is not None
        rootNode = self.parent is None
        assert(self.values is not None)
        # um noh interno(com excessao da raiz) tem, ao menos, min_values
        if not rootNode and innerNode:
            assert(tree.min_values <= len(self.values))
        # um noh nao pode ter mais de max_values
        assert(len(self.values) <= tree.max_values)
        # a raiz tem, ao menos, dois filhos se nao for uma folha.
        if rootNode and innerNode:
            assert(len(self.children) >= 2)
        # Um noh que nao seja uma folha com k filhos contem k-1 chaves.
        if innerNode:
            assert((len(self.values) + 1) == len(self.children))
        # checa se os valores estao ordenados
        prev = None
        for i in self.values:
            if prev is not None:
                assert(i > prev)
            prev = i
        if self.children:
            for i in self.children:
                assert(i.parent is self)
                i.check_valid(tree)

    def search(self, val):
        '''
        retorna uma tupla. Se o valor existir (True, noh, posicao),
        caso contrario (False, noh, posicao onde deveria ser inserido)
        '''
        i = bisect.bisect_left(self.values, val)
        if (i != len(self.values) and not val < self.values[i]):
            # o valor foi encontrado
            assert(self.values[i] == val)
            return (True, self, i)

        if self.children is not None:
            assert(len(self.children) >= i and self.children[i])
            # busca recursivamente o noh filho apropriado
            return self.children[i].search(val)
        else:
            return (False, self, i)

    def _split_node(self, tree, val=None, slot=None, childNodes=None):
        '''
        quebra uma arvore B em duas. Se val for fornecido, insere-o nos nohs resultantes.
        '''
        assert(val is None or (slot is not None))
        midList = [] if val is None else [ val ]
        if slot is None:
            slot = 0
        # pega a media de self.values e val
        splitValues = self.values[0:slot] + midList + self.values[slot:]
        medianIdx = len(splitValues) // 2
        lv = splitValues[0:medianIdx]
        medianVal = splitValues[medianIdx]
        rv = splitValues[medianIdx + 1:]
        innerNode = self.children is not None
        if innerNode:
            if childNodes is not None:
                splitChildren = (self.children[0:slot] +
                                 list(childNodes) +
                                 self.children[slot + 1:])
            else:
                splitChildren = self.children
            lc = splitChildren[0:len(lv) + 1]
            rc = splitChildren[len(lv) + 1:]
        else:
            lc = None
            rc = None

        leftNode = _BTreeNode(lv, lc)
        rightNode = _BTreeNode(rv, rc)
        if self.parent:
            self.parent.add(tree,
                            medianVal,
                            None,
                            (leftNode, rightNode))
        else:
            # cria nova raiz e incrementa a profundidade da arvore
            newRoot = _BTreeNode([ medianVal ], [leftNode, rightNode])
            leftNode.parent = newRoot
            rightNode.parent = newRoot
            tree.root = newRoot
            tree.height += 1
            tree.size += 1

    def add(self, tree, val, slot=None, childNodes=None):
        '''
        adiciona um novo valor na arvore. o valor nao pode existir previamente
        '''
        assert(self.children is None or childNodes)
        innerNode = self.children is not None
        if innerNode:
            assert(childNodes and len(childNodes) == 2)
        else:
            assert(childNodes is None)

        if slot is None:
            slot = bisect.bisect_left(self.values, val)

        if len(self.values) < tree.max_values:
            self.values.insert(slot, val)
            tree.size += 1
            if childNodes:
                for i in childNodes:
                    i.parent = self
                self.children[slot:slot + 1] = childNodes
            return True
        self._split_node(tree, val, slot, childNodes)
        return True


    def min_value(self, slot=0):
        if self.children:
            return self.children[slot].min_value()
        return self.values[0], self, 0

    def max_value(self, slot=None):
        if slot is None:
            slot = len(self.values) - 1
        if self.children:
            return self.children[slot + 1].max_value()
        return self.values[-1], self, len(self.values) - 1


    def delete(self, tree, val, slot=None):
        '''
        deleta um valor da arvore. o valor tem de existir previamente
        '''
        innerNode = self.children is not None
        if slot is None:
            assert(slot is not None)
            slot = bisect.bisect_left(self.values, val)
        assert(slot != len(self.values) and self.values[slot] == val)
        if not innerNode:
            del self.values[slot]
            tree.size -= 1
            if len(self.values) < tree.min_values:
                self._rebalance(tree)
        else:
            newSep, node, idx = self.min_value(slot + 1)
            self.values[slot] = newSep
            del node.values[idx]
            tree.size -= 1
            if len(node.values) < tree.min_values:
                node._rebalance(tree)

    def _rebalance(self, tree):
        '''
        rebalanceia a arvore comecando no noh atual
        '''
        lsibling, rsibling, idx = self.get_siblings()
        assert(rsibling or lsibling or self.parent is None)
        if self.parent is None:
            return
        innerNode = self.children is not None
        if innerNode:
            assert(rsibling is None or rsibling.children is not None)
            assert(lsibling is None or lsibling.children is not None)
        else:
            assert(rsibling is None or rsibling.children is None)
            assert(lsibling is None or lsibling.children is None)
        if not innerNode:
            if rsibling and len(rsibling.values) > tree.min_values:
                sepIdx = idx
                sepVal = self.parent.values[sepIdx]
                self.parent.values[sepIdx] = rsibling.values[0]
                del rsibling.values[0]
                self.values.append(sepVal)
                return
            elif lsibling and len(lsibling.values) > tree.min_values:
                sepIdx = idx - 1
                sepVal = self.parent.values[sepIdx]
                self.parent.values[sepIdx] = lsibling.values[-1]
                del lsibling.values[-1]
                self.values.insert(0, sepVal)
                return

        if lsibling is not None:
            sepIdx = idx - 1
            ln = lsibling
            rn = self
        elif rsibling is not None:
            sepIdx = idx
            ln = self
            rn = rsibling
        else:
            assert(False)

        sepVal = self.parent.values[sepIdx]

        ln.values.append(sepVal)
        ln.values.extend(rn.values)
        del rn.values[:]
        del self.parent.values[sepIdx]
        assert(self.parent.children[sepIdx + 1] is rn)
        del self.parent.children[sepIdx + 1]
        if rn.children:
            ln.children.extend(rn.children)
            for i in rn.children:
                i.parent = ln

        if len(ln.values) > tree.max_values:
            # we have to split the newly formed node
            # this situation can aris only when merging inner nodes
            assert(innerNode)
            ln._split_node(tree)

        if len(self.parent.values) < tree.min_values:
            # rebalance the parent
            self.parent._rebalance(tree)

        if self.parent.parent is None and not self.parent.values:
            tree.root = ln
            tree.root.parent = None

    def get_siblings(self):
        if not self.parent:
            # a raiz nao tem irmaos
            return (None, None, 0)
        assert(self.parent.children)
        lsibling = None
        rsibling = None
        idx = 0
        for i, j in enumerate(self.parent.children):
            if j is self:
                if i != 0:
                    lsibling = self.parent.children[i - 1]
                if (i + 1) < len(self.parent.children):
                    rsibling = self.parent.children[i + 1]
                idx = i
                break
        return (lsibling, rsibling, idx)

class BTree(object):
    '''
    cria uma arvore b de ordem 'order'(numero maximo de filhos por no'),
    que e' o numero maximo de chaves por no' + 1. Retirado de
    The Art of Computer Programming, Knuth, Volume 3 p. 483.
    '''
    def __init__(self, order):
        if order <= 2:
            raise ValueError("a ordem da arvore b deve ser, ao menos, 3")
        self.root = _BTreeNode()
        self.order = order
        self.max_values = order - 1
        self.min_values = self.max_values // 2
        self.height = 1
        self.size = 0

    def __str__(self):
        return 'altura: %d itens: %d m: %d raiz: %x' % (
                                    self.height, self.size,
                                    self.max_values + 1,
                                    id(self.root))

    def add(self, val):
        # encontra a folha onde o valor tem de ser inserido
        found, node, slot = self.root.search(val)
        if found:
            # o valor ja existe, nao faz nada
            return False
        return node.add(self, val, slot, None)

    def delete(self, val):
        # encontra o valor
        found, node, slot = self.root.search(val)
        if not found:
            # o valor nao existe, nao faz nada
            return False
        return node.delete(self, val, slot)

    def search(self, val):
        return self.root.search(val)[0]

    def min(self):
        return self.root.min_value()[0]

    def max(self):
        return self.root.max_value()[0]


if __name__ == '__main__':
    # mini teste
    tree = BTree(3)
    for i in range(20):
        tree.add(random.randint(0,1000))
        #assert(tree.search(i))
    #for i in range(1, 8):
    #    assert(tree.search(i))

    print("arvore-b")
    tree.root.pretty_print()
    print("\n-*o* - - - - - - - - - - - - - - - - *o*-")
    tree.root.check_valid(tree)

    for i in range(1, 8):
        tree.delete(i)
        tree.root.check_valid(tree)