day 8: cleaner

1 files changed, 9 insertions, 20 deletions

diff --git a/8.py b/8.py
index 0bd67ed..609d09f 100644
--- a/8.py
+++ b/8.py
@@ -26,16 +26,7 @@ dig = [
     fset((Seg.A, Seg.B, Seg.C, Seg.D, Seg.E, Seg.F, Seg.G)),
     fset((Seg.A, Seg.B, Seg.C, Seg.D, Seg.F, Seg.G)),
 ]
-
-seg = [
-    fset((0, 2, 3, 5, 6, 7, 8, 9)),
-    fset((0, 4, 5, 6, 8, 9)),
-    fset((0, 1, 2, 3, 4, 7, 8, 9)),
-    fset((2, 3, 4, 5, 6, 8, 9)),
-    fset((0, 2, 6, 8)),
-    fset((0, 1, 3, 4, 5, 6, 7, 8, 9)),
-    fset((0, 2, 3, 5, 6, 8, 9)),
-]
+seg = [ fset(i for i in range(10) if seg in dig[i]) for seg in Seg ]
 
 counts = Counter()
 count = 0
@@ -47,24 +38,22 @@ with open_day(8) as f:
             count: seg
             for seg, count in Counter(chain.from_iterable(patterns)).items()
         }
-        def patbylen(l): return next(p for p in patterns if len(p) == l)
-        CF      = patbylen(len(dig[1]))
-        BCDF    = patbylen(len(dig[4]))
-        ACF     = patbylen(len(dig[7]))
-        ABCDEFG = patbylen(len(dig[8]))
+        lenmap = { len(pat): pat for pat in patterns }
+        CF      = lenmap[len(dig[1])]
+        BCDF    = lenmap[len(dig[4])]
+        ACF     = lenmap[len(dig[7])]
+        ABCDEFG = lenmap[len(dig[8])]
         segmap = dict()
         segmap[Seg.B] = countmap[len(seg[Seg.B])]
         segmap[Seg.E] = countmap[len(seg[Seg.E])]
         segmap[Seg.F] = countmap[len(seg[Seg.F])]
-        def single(s):
-            assert(len(s) == 1)
-            return next(iter(s))
+        def single(s): return next(iter(s))
         segmap[Seg.A] = single(ACF - CF)
         segmap[Seg.C] = single(CF - fset(segmap[Seg.F]))
         segmap[Seg.D] = single(BCDF - CF - fset(segmap[Seg.B]))
         segmap[Seg.G] = single(ABCDEFG - fset(segmap.values()))
-        digmap = { fset(segmap[s] for s in dig[i]): i for i in range(10) }
-        digits = [digmap[fset(d)] for d in digits.split()]
+        digmap = { fset(segmap[s] for s in d): i for i, d in enumerate(dig) }
+        digits = [ digmap[fset(d)] for d in digits.split() ]
         counts.update(digits)
         count += sum(d * 10 ** i for i, d in enumerate(reversed(digits)))
 print(counts[1] + counts[4] + counts[7] + counts[8])