Well, Well, Well... 2018/03/02 code puzzles jane-street Fun puzzle, part of the resolve series. import collections import doctest from fractions import Fraction def label2Di(graph): ''' >>> G = graph2Di(""" ... 9 9 9 12 ... 9 73 73 12 ... 9 9 73 12 ... 73 73 73 73""") >>> L = label2Di(G) >>> print(string2Di({k: v for v, group in enumerate(L) for k in group})) 0 0 0 1 0 2 2 1 0 0 2 1 2 2 2 2 ''' L = [] while graph: L.append(set()) available = {min(graph, key=list(graph).index)} while available: k = available.pop() v = graph.pop(k) L[-1].add(k) available.update(k+s for s in (1, -1, 1j, -1j) if graph.get(k+s) == v) return L def graph2Di(string): return { x+1j*y: int(v) if v.isdigit() else v for y, row in enumerate(string.strip().splitlines()) for x, v in enumerate(row.split()) } def string2Di(graph): X, Y = max(int(p.real)+1 for p in graph), max(int(p.imag)+1 for p in graph) template = "\n".join("".join(f"{{{x+1j*y}:>3}}" for x in range(X)) for y in range(Y)) return template.format_map({str(k): int(v) for k, v in graph.items()}) def resolve(D): # analyze topology A = {k: group for group in label2Di({**D}) for k in group} # mesas F = {p: {p+s for s in (1, -1, 1j, -1j) if D.get(p+s, 0) > D[p]} for p in D} # flows (downstream) # distribute sources = collections.Counter({0: Fraction(1)}) sinks = collections.Counter() while sources: copy = {**sources} sources.clear() for p, q in {**copy}.items(): flows = collections.Counter(f for a in A[p] for f in F[a]) if flows: sources += {k: q * v / sum(flows.values()) for k, v in flows.items()} else: r = min(A[p], key=str) sinks[r] += q # determine time-step and fill up areas = {p: len(A[p]) for p in sinks} heights = {p: min(D[p] - D[m] for m in F if A[p] & F[m]) for p in sinks} times = {p: (areas[p] * heights[p]) / f for p, f in sinks.items()} k = min(times, key=times.get) D -= {a: f * times[k] / areas[p] for p, f in sinks.items() for a in A[p]} return times[k] import doctest doctest.testmod() string = ''' 1 5 27 22 28 40 14 39 13 17 30 41 12 2 32 35 24 25 19 47 34 16 33 10 42 7 44 18 3 8 45 37 4 21 20 15 46 38 6 26 48 49 9 23 31 29 11 36 43''' D = collections.Counter({k: int(v) for k, v in graph2Di(string).items()}) print(sum(resolve(D) for _ in range(25))) print(string2Di(D))