# treeset.py
# Ryosuke Fukatani, 2016 (modified)
# https://github.com/fukatani/TreeSet


import bisect

class TreeSet(object):
  '''
    Binary-tree set like java Treeset.
    Duplicate elems will not be added.
    When added new elem, TreeSet will be sorted automatically.
  '''

  def __init__(self):
    self.tset = []
    self.popNum = 0
    self.pushNum = 0


  def isEmpty(self):
    return len(self) == 0


  def addAll(self, elems):
    for elem in elems:
      if elem in self: 
        continue
      self.add(elem)


  def add(self,elem):
    if self.isEmpty() or elem not in self:
      bisect.insort_right(self.tset, elem)
      self.pushNum += 1


  def addHighLow(self, elem):
    if elem not in self:
      idx = bisect.bisect_right(self.tset, elem)
      self.tset.insert(idx, elem)

      low = self.tset[idx-1] if idx > 0 else None
      high = self.tset[idx+1] if idx < len(self.tset) - 1 else None
      self.pushNum += 1
      return low, high
    return None


  def push(self, elem):
    self.add(elem)


  def pushAll(self, elems):
    for elem in elems:
      self.push(elem)


  def pop(self):
    if not self.isEmpty():
      res = self.tset[0]
      self.tset = self.tset[1:]
      return res
    return None


  # Returns the least elem in this set strictly greater than the given elem,
  # or None if there is no such elem.
  def higher(self, e):
    idx = bisect.bisect_right(self.tset, e)
    if idx < self.__len__():
      return self[idx]
    else:
      return None


  # Returns the greatest elem in this set strictly less than the given elem,
  # or None if there is no such elem.
  def lower(self, e):
    idx = bisect.bisect_left(self.tset, e)
    if idx > 0:
      return self[idx - 1]
    else:
      return None


  def swap(self, elem1, elem2):
    idx1 = self.tset.index(elem1)
    idx2 = self.tset.index(elem2)
    self.tset[idx1] = elem2
    self.tset[idx2] = elem1


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


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


  def clear(self):
    self.tset = []


  def remove(self, elem):
    try:
      self.tset.remove(elem)
    except ValueError:
      return False
    return True


  def __iter__(self):
    for elem in self.tset:      # ascending iteration
      yield elem


  def __str__(self):
    return '[' + ' '.join([str(i) for i in self.tset]) + ']'


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


  def __contains__(self, e):
    # Fast attribution judgment by bisect
    try:
      return e == (self.tset[bisect.bisect_left(self.tset, e)])
    except:
      return False


# ------------------------------------------

if __name__ == '__main__':
  ts = TreeSet()
  ts.addAll([3,7,7,1,3])
  print("Tree set:", ts)
  print("> 4:", ts.higher(4))
  print("< 4:", ts.lower(4))
  print("> 3:", ts.higher(3))
  print("< 3:", ts.lower(3))
  print("< 1:", ts.lower(1))
  print("> 7:", ts.higher(7))