96 lines
2.7 KiB
Python
96 lines
2.7 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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|