def test_element_not_present():
elements = get_elements(n=10)
dis = DisjointSet(elements)
with assert_raises(KeyError):
dis["dummy"]
with assert_raises(KeyError):
dis.merge(elements[0], "dummy")
with assert_raises(KeyError):
dis.connected(elements[0], "dummy")
def test_self_unions(n):
elements = get_elements(n)
dis = DisjointSet(elements)
for x in elements:
assert dis.connected(x, x)
assert not dis.merge(x, x)
assert dis.connected(x, x)
assert dis.n_subsets == len(elements)
assert elements == list(dis)
roots = [dis[x] for x in elements]
assert elements == roots
def test_binary_tree(kmax):
n = 2**kmax
elements = get_elements(n)
dis = DisjointSet(elements)
rng = np.random.RandomState(seed=0)
for k in 2**np.arange(kmax):
for i in range(0, n, 2 * k):
r1, r2 = rng.randint(0, k, size=2)
a, b = elements[i + r1], elements[i + k + r2]
assert not dis.connected(a, b)
assert dis.merge(a, b)
assert dis.connected(a, b)
assert elements == list(dis)
roots = [dis[i] for i in elements]
expected_indices = np.arange(n) - np.arange(n) % (2 * k)
expected = [elements[i] for i in expected_indices]
assert roots == expected
def test_linear_union_sequence(n, direction):
elements = get_elements(n)
dis = DisjointSet(elements)
assert elements == list(dis)
indices = list(range(n - 1))
if direction == "backwards":
indices = indices[::-1]
for it, i in enumerate(indices):
assert not dis.connected(elements[i], elements[i + 1])
assert dis.merge(elements[i], elements[i + 1])
assert dis.connected(elements[i], elements[i + 1])
assert dis.n_subsets == n - 1 - it
roots = [dis[i] for i in elements]
if direction == "forwards":
assert all(elements[0] == r for r in roots)
else:
assert all(elements[-2] == r for r in roots)
assert not dis.merge(elements[0], elements[-1])