118 lines
3.4 KiB
Python
118 lines
3.4 KiB
Python
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import unittest
|
|
from dsu_rollback import DSU
|
|
|
|
|
|
class TestDSURollback(unittest.TestCase):
|
|
def test_basic_union_connected(self):
|
|
d = DSU(4)
|
|
self.assertTrue(d.union(0, 1))
|
|
self.assertTrue(d.connected(0, 1))
|
|
self.assertFalse(d.connected(0, 2))
|
|
self.assertEqual(d.size(0), 2)
|
|
|
|
def test_noop_union_returns_false_no_history(self):
|
|
d = DSU(3)
|
|
d.union(0, 1)
|
|
m = d.snapshot()
|
|
self.assertFalse(d.union(0, 1))
|
|
self.assertEqual(d.snapshot(), m)
|
|
|
|
def test_rollback_restores(self):
|
|
d = DSU(5)
|
|
d.union(0, 1); d.union(2, 3)
|
|
m = d.snapshot()
|
|
d.union(1, 2); d.union(3, 4)
|
|
self.assertTrue(d.connected(0, 4))
|
|
d.rollback(m)
|
|
self.assertFalse(d.connected(0, 4))
|
|
self.assertTrue(d.connected(0, 1))
|
|
self.assertTrue(d.connected(2, 3))
|
|
self.assertFalse(d.connected(1, 2))
|
|
|
|
def test_nested_snapshots(self):
|
|
d = DSU(6)
|
|
d.union(0, 1)
|
|
m1 = d.snapshot()
|
|
d.union(2, 3)
|
|
m2 = d.snapshot()
|
|
d.union(1, 2)
|
|
d.rollback(m2)
|
|
self.assertTrue(d.connected(0, 1) and d.connected(2, 3))
|
|
self.assertFalse(d.connected(0, 2))
|
|
d.rollback(m1)
|
|
self.assertTrue(d.connected(0, 1))
|
|
self.assertFalse(d.connected(2, 3))
|
|
|
|
def test_full_rollback_to_zero(self):
|
|
d = DSU(4)
|
|
d.union(0, 1); d.union(2, 3); d.union(1, 2)
|
|
d.rollback(0)
|
|
for i in range(4):
|
|
self.assertEqual(d.size(i), 1)
|
|
self.assertEqual(d.components(), 4)
|
|
|
|
def test_invalid_marker(self):
|
|
d = DSU(2)
|
|
with self.assertRaises(ValueError):
|
|
d.rollback(5)
|
|
with self.assertRaises(ValueError):
|
|
d.rollback(-1)
|
|
|
|
def test_no_path_compression_find_stable(self):
|
|
d = DSU(4)
|
|
d.union(0, 1); d.union(1, 2)
|
|
before = list(d.parent)
|
|
r = d.find(2)
|
|
self.assertEqual(r, d.find(0))
|
|
self.assertEqual(d.parent[2], before[2])
|
|
|
|
def test_size_after_rollback(self):
|
|
d = DSU(3)
|
|
d.union(0, 1)
|
|
m = d.snapshot()
|
|
d.union(0, 2)
|
|
self.assertEqual(d.size(0), 3)
|
|
d.rollback(m)
|
|
self.assertEqual(d.size(0), 2)
|
|
self.assertEqual(d.size(2), 1)
|
|
|
|
def test_components(self):
|
|
d = DSU(4)
|
|
self.assertEqual(d.components(), 4)
|
|
d.union(0, 1); d.union(2, 3)
|
|
self.assertEqual(d.components(), 2)
|
|
d.union(1, 2)
|
|
self.assertEqual(d.components(), 1)
|
|
|
|
def test_negative_n(self):
|
|
with self.assertRaises(ValueError):
|
|
DSU(-1)
|
|
|
|
|
|
def test_shape_union_by_size(self):
|
|
d = DSU(8)
|
|
d.union(0,1); d.union(2,3); d.union(4,5); d.union(6,7)
|
|
d.union(1,2); d.union(4,0)
|
|
# union by size: attach smaller root under larger -> parent[4]==0, depth 2
|
|
self.assertEqual(d.parent[4], 0)
|
|
# deepest depth
|
|
maxdepth = 0
|
|
for i in range(8):
|
|
depth=0; cur=i
|
|
while d.parent[cur]!=cur:
|
|
cur=d.parent[cur]; depth+=1
|
|
maxdepth=max(maxdepth,depth)
|
|
self.assertEqual(maxdepth, 2)
|
|
|
|
def test_rollback_marker_type(self):
|
|
d = DSU(2)
|
|
for bad in (True, False, 1.5, "x"):
|
|
with self.assertRaises((TypeError, ValueError)):
|
|
d.rollback(bad)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|