def test_treeset(self):
ts = TreeSet([3,7,2,7,1,3])
self.assertEqual(ts, [1, 2, 3, 7])
ts.add(4)
self.assertEqual(ts, [1, 2, 3, 4, 7])
ts.add(4)
self.assertEqual(ts, [1, 2, 3, 4, 7])
ts.remove(7)
self.assertEqual(ts, [1, 2, 3, 4])
ts.remove(5)
self.assertEqual(ts, [1, 2, 3, 4])
ts.addAll([3,4,5,6])
self.assertEqual(ts, [1, 2, 3, 4, 5, 6])
ts.pop(3)
self.assertEqual(ts, [1, 2, 3, 5, 6])
self.assertEqual(ts[0], 1)
self.assertEqual(ts[-1], 6)
self.assertTrue(1 in ts)
self.assertFalse(100 in ts)
for i, element in enumerate(TreeSet([1,3,1])):
if i==0:
self.assertEqual(element, 1)
elif i==1:
self.assertEqual(element, 3)
else:
raise Exception
ts_copy = ts.clone()
self.assertEqual(ts, [1, 2, 3, 5, 6])
self.assertEqual(ts.floor(4), 3)
self.assertEqual(ts.ceiling(4), 5)
self.assertEqual(ts.floor(3), 3)
self.assertEqual(ts.ceiling(3), 3)
ts.clear()
self.assertEqual(ts._treeset, [])
def test_treeset(self):
ts = TreeSet([3, 7, 2, 7, 1, 3])
self.assertEqual(ts, [1, 2, 3, 7])
ts.add(4)
self.assertEqual(ts, [1, 2, 3, 4, 7])
ts.add(4)
self.assertEqual(ts, [1, 2, 3, 4, 7])
ts.remove(7)
self.assertEqual(ts, [1, 2, 3, 4])
ts.remove(5)
self.assertEqual(ts, [1, 2, 3, 4])
ts.addAll([3, 4, 5, 6])
self.assertEqual(ts, [1, 2, 3, 4, 5, 6])
ts.pop(3)
self.assertEqual(ts, [1, 2, 3, 5, 6])
self.assertEqual(ts[0], 1)
self.assertEqual(ts[-1], 6)
self.assertTrue(1 in ts)
self.assertFalse(100 in ts)
for i, element in enumerate(TreeSet([1, 3, 1])):
if i == 0:
self.assertEqual(element, 1)
elif i == 1:
self.assertEqual(element, 3)
else:
raise Exception
ts_copy = ts.clone()
self.assertEqual(ts, [1, 2, 3, 5, 6])
self.assertEqual(ts.floor(4), 3)
self.assertEqual(ts.ceiling(4), 5)
self.assertEqual(ts.floor(3), 3)
self.assertEqual(ts.ceiling(3), 3)
ts.clear()
self.assertEqual(ts._treeset, [])
class TreeMap(dict):
"""
"TreeMap" is a dictionary with sorted keys similar to java TreeMap.
Keys, iteration, items, values will all return values ordered by key.
Otherwise it should behave just like the builtin dict.
"""
def __init__(self, seq=None, **kwargs):
if seq is None:
super().__init__(**kwargs)
else:
super().__init__(seq, **kwargs)
self.sorted_keys = TreeSet(super().keys())
def __setitem__(self, key, value):
super().__setitem__(key, value)
self.sorted_keys.add(key)
def __delitem__(self, key):
super().__delitem__(key)
self.sorted_keys.remove(key)
def keys(self):
return self.sorted_keys
def items(self):
return [(k, self[k]) for k in self.sorted_keys]
def __iter__(self):
for k in self.sorted_keys:
yield k
def values(self):
for k in self.sorted_keys:
yield self[k]
def clear(self):
super().clear()
self.sorted_keys.clear()