3b

three.py

@@ -5,14 +5,26 @@ from useful import *
 text = open(0).read()
 w, h = re.search('\n', text).start(), len(re.findall('\n', text))
 symmap = np.zeros((h, w), dtype=bool)
+gearmap = np.zeros((h, w))
+geardict = defaultdict(lambda: [])
 
 def numcount(map, match):
     beg, end = ((c // (w+1), c % (w+1)) for c in (match.start(), match.end()))
     return int(match.group(0)) if any(map[beg[0],beg[1]:end[1]]) else 0
 
-def symput(map, match):
-    y, x = (match.start() // (w+1), match.start() % (w+1))
-    map[y-1:y+2, x-1:x+2] = True
+def gearnumcount(map, match):
+    beg, end = ((c // (w+1), c % (w+1)) for c in (match.start(), match.end()))
+    for c in map[beg[0],beg[1]:end[1]]:
+        if c:
+            geardict[c].append(int(match.group(0)))
+            return # assume each number touches just _one_ gear
 
-_ = [symput(symmap, m) for m in re.finditer(r'[^\n0-9.]', text)]
+def symput(map, match, val):
+    y, x = (match.start() // (w+1), match.start() % (w+1))
+    map[y-1:y+2, x-1:x+2] = val
+
+_ = [symput(symmap, m, True) for m in re.finditer(r'[^\n0-9.]', text)]
+_ = [symput(gearmap, m, i + 1) for i, m in enumerate(re.finditer(r'[*]', text))]
 print(sum(numcount(symmap, m) for m in re.finditer(r'[1-9][0-9]*', text)))
+_ = [gearnumcount(gearmap, m) for m in re.finditer(r'[1-9][0-9]*', text)]
+print(sum(n[0] * n[1] for n in geardict.values() if len(n) == 2))