Python Library for Self-Balancing BST

Here we will see simulating the library framework TreeSet which is available in Java on Python. There are situations that arise to disallow duplicate entries especially in any management software where an object is uniquely identified by its uniqueness. In this post, let us see how to do that in Pythonic way.

In our implementation, “TreeSet” class is a Binary-tree set like the Java TreeSet. The TreeSet will be sorted automatically when adding/removing elements. Duplicate elements will not be added.

The functionalities that have been included are :

• Inserting an element into the TreeSet.
• Inserting multiple elements into the TreeSet.
• Fetching the ceil value from the TreeSet.
• Fetching the floor value from the TreeSet.
• Deleting an element from the TreeSet.
• Deleting all the elements from the TreeSet.
• Fetching the shallow copy of the TreeSet.
• Popping an element from the TreeSet.

Python3

# The bisect module ensures the automatic sort after insertion
import bisect

class TreeSet(object):
    """
    Binary-tree set like java Treeset.
    Duplicate elements will not be added.
    When added new element, TreeSet will be
    sorted automatically.
    """
    def __init__(self, elements):
        self._treeset = []
        self.addAllElements(elements)

    # To add many elements
    def addAllElements(self, elements):
        for element in elements:
            if element in self: continue
            self.addElement(element)

    # To add an element
    def addElement(self, element):
        if element not in self:
            bisect.insort(self._treeset, element)

    # To get ceiling value
    def ceiling(self, e):
        index = bisect.bisect_right(self._treeset, e)
        if self[index - 1] == e:
            return e
        return self._treeset[bisect.bisect_right(self._treeset, e)]

    def floor(self, e):
        index = bisect.bisect_left(self._treeset, e)
        if self[index] == e:
            return e
        else:
            return self._treeset[bisect.bisect_left(self._treeset, e) - 1]

    def __getitem__(self, num):
        return self._treeset[num]

    def __len__(self):
        return len(self._treeset)

    def clearElements(self):
        """
        Delete all elements in TreeSet.
        """
        self._treeset = []

    def clone(self):
        """
        Return shallow copy of self.
        """
        return TreeSet(self._treeset)

    def removeElement(self, element):
        """
        Remove element if element in TreeSet.
        """
        try:
            self._treeset.remove(element)
        except ValueError:
            return False
        return True

    def __iter__(self):
        """
        Do ascending iteration for TreeSet
        """
        for element in self._treeset:
            yield element

    def pop(self, index):
        return self._treeset.pop(index)

    def __str__(self):
        return str(self._treeset)

    def __eq__(self, target):
        if isinstance(target, TreeSet):
            return self._treeset == target.treeset
        elif isinstance(target, list):
            return self._treeset == target

    def __contains__(self, e):
        """
        Fast attribution judgement by bisect
        """
        try:
            return e == self._treeset[bisect.bisect_left(self._treeset, e)]
        except:
            return False

if __name__ == '__main__':
    treeSet = TreeSet([3, 7, 7, 1, 3, 10])

    # As there is no 5, floor of this gives the
    # next least value  in the treeset i.e. 3
    print("treeSet.floor(5) : ", treeSet.floor(5))

    # As there is no 4, ceil of this gives the
    # next highest value in the treeset i.e. 7
    print("treeSet.ceiling(4) : ", treeSet.ceiling(4))

    # As there is no 2, floor of this gives the next
    # least value  in the treeset i.e. 1
    print("treeSet.floor(2) : ", treeSet.floor(2))

    # As there is 3, ceil will give 3 as it is
    print("treeSet.ceiling(3) : ", treeSet.ceiling(3))

    # As there is 10, floor will give 10 as it is
    print("treeSet.floor(10) : ", treeSet.floor(10))

    # No duplicates printed
    print("treeSet : ", treeSet)

    # 2nd example
    treeSet = TreeSet([30, 70, 20, 70, 10, 30])

    # No duplicates printed
    print("2nd example treeSet :", treeSet)

    treeSet.addElement(40)
    print("Adding 40 to the above, it is arranged in\
    sorted order : ", treeSet)

    treeSet.removeElement(70)
    print("After removing 70 ", treeSet)

    treeSet.removeElement(50)
    print("After removing 50, no issue even if not\
    available ", treeSet)
    # There is no 50 available

    treeSet.addAllElements([30, 40, 50, 60])
    print("After adding many elements, duplicate are\
    not taken :", treeSet)

    # Getting first element of treeset
    print(treeSet[0])

    # See, how elements can be retrieved
    print(treeSet[-1], treeSet[5], treeSet[-2],
          treeSet[-3], treeSet[-4], treeSet[-5])

    # Check for the existence of 10 and since
    # found, it returns true
    print(10 in treeSet)

    # Check for the existence of 100 and since
    # not found, it returns false
    print(100 in treeSet)

    for i in TreeSet([10, 30, 40]):
        print(i)

Output :

treeSet.floor(5) : 3
treeSet.ceiling(4) : 7
treeSet.floor(2) : 1
treeSet.ceiling(3) : 3
treeSet.floor(10) : 10
treeSet : [1, 3, 7, 10]
treeSet : [10, 20, 30, 70]
Adding 40 to the above, it is arranged in sorted order : [10, 20, 30, 40, 70]
After removing 70 [10, 20, 30, 40]
After removing 50 [10, 20, 30, 40]
After adding many elements : [10, 20, 30, 40, 50, 60]
10
60 60 50 40 30 20
True
False
10
30
40