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)