#!/usr/bin/env python3 import re lines = (x.strip() for x in open("input.txt")) found = set() nums = [dict(), dict(), dict()] signals = [] DIGIT = re.compile("\d+") for idx, line in enumerate(lines): next_num_items = [(idx, match.start(), match.end(), int(line[match.start():match.end()])) for match in DIGIT.finditer(line)] next_num_items = {i: val for val in next_num_items for i in range(val[1], val[2])} next_char_idxs = [j for j, char in enumerate(line) if char not in "0123456789."] nums = nums[1:] + [next_num_items] # check for signal in signals: for i in range(signal-1, signal+2): for j in range(0, 3): if i in nums[j]: found.add(nums[j][i]) signals = next_char_idxs print(sum(map(lambda x: x[3], found)))