diff --git a/test_treeset.py b/test_treeset.py index da42fae..86c3071 100644 --- a/test_treeset.py +++ b/test_treeset.py @@ -50,6 +50,10 @@ def test_treeset(self): ts.clear() self.assertEqual(ts._treeset, []) + + ts = TreeSet([int(x) for x in '934853458236']) + self.assertEqual(ts.ceiling(100), None) + self.assertEqual(ts.floor(0), None) if __name__ == '__main__': unittest.main() diff --git a/treeset.py b/treeset.py index f601f9b..1cf1b3e 100644 --- a/treeset.py +++ b/treeset.py @@ -32,14 +32,19 @@ 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)] + try: + return self._treeset[bisect.bisect_right(self._treeset, e)] + except IndexError: + return None 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] + check = self._treeset[bisect.bisect_left(self._treeset, e) - 1] + if check <= e: + return check + return None def __getitem__(self, num): return self._treeset[num] @@ -48,21 +53,15 @@ def __len__(self): return len(self._treeset) def clear(self): - """ - Delete all elements in TreeSet. - """ + """Delete all elements in TreeSet.""" self._treeset = [] def clone(self): - """ - Return shallow copy of self. - """ + """Return shallow copy of self.""" return TreeSet(self._treeset) def remove(self, element): - """ - Remove element if element in TreeSet. - """ + """Remove element if element in TreeSet.""" try: self._treeset.remove(element) except ValueError: @@ -70,9 +69,7 @@ def remove(self, element): return True def __iter__(self): - """ - Do ascending iteration for TreeSet - """ + """Do ascending iteration for TreeSet""" for element in self._treeset: yield element @@ -87,14 +84,13 @@ def __eq__(self, target): return self._treeset == target.treeset elif isinstance(target, list): return self._treeset == target + return None def __contains__(self, e): - """ - Fast attribution judgment by bisect - """ + """Fast attribution judgment by bisect""" try: return e == self._treeset[bisect.bisect_left(self._treeset, e)] - except: + except Exception: return False if __name__ == '__main__':