def day9b(s): from scipy.cluster.hierarchy import DisjointSet grid = np.array([[int(a) for a in line] for line in s.splitlines()], dtype=int) ymax, xmax = grid.shape loc = [(y, x) for y in range(ymax) for x in range(xmax) if grid[y, x] < 9] basins = DisjointSet(loc) for y, x in loc: neighbors = [(grid[y + yd, x + xd], y + yd, x + xd) for yd, xd in ((-1, 0), (0, 1), (0, -1), (1, 0)) if 0 <= y + yd < ymax and 0 <= x + xd < xmax and grid[y, x] < 9] val, yy, xx = min(neighbors, default=(999, -1, -1)) if grid[y, x] > val: basins.merge((y, x), (yy, xx)) a, b, c = sorted(basins.subsets(), key=len)[-3:] return len(a) * len(b) * len(c)
def test_subsets(n):
elements = get_elements(n)
dis = DisjointSet(elements)
rng = np.random.RandomState(seed=0)
for i, j in rng.randint(0, n, (n, 2)):
x = elements[i]
y = elements[j]
expected = {element for element in dis if {dis[element]} == {dis[x]}}
assert expected == dis.subset(x)
expected = {dis[element]: set() for element in dis}
for element in dis:
expected[dis[element]].add(element)
expected = list(expected.values())
assert expected == dis.subsets()
dis.merge(x, y)
assert dis.subset(x) == dis.subset(y)
