Related Articles

Related Articles

Python Library for Self-Balancing BST
  • Last Updated : 18 Aug, 2020

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.
filter_none

edit
close

play_arrow

link
brightness_4
code

# 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 judgment 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)   

chevron_right


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

Attention geek! Strengthen your foundations with the Python Programming Foundation Course and learn the basics.

To begin with, your interview preparations Enhance your Data Structures concepts with the Python DS Course.




My Personal Notes arrow_drop_up
Recommended Articles
Page :