diff --git a/3/engine.py b/3/engine.py index 6f19640..b282ab2 100644 --- a/3/engine.py +++ b/3/engine.py @@ -18,8 +18,8 @@ def parse_input(input_data: str) -> dict: # convert input into a 2d array for easier parsing field = [line for line in input_data.split("\n")] - # get coordinates of numbers - nums = { + # create a dictionary of all symbols and their coordinates + symbols = { (i, char): [] for i in range(140) for char in range(140) @@ -28,16 +28,18 @@ def parse_input(input_data: str) -> dict: # for each row in field for i, row in enumerate(field): - # for each number in row + # for each number in row create a dictionary of any number with coords for n in re.finditer(r"\d+", row): - edge = { + numbers = { (i, char) for i in (i - 1, i, i + 1) for char in range(n.start() - 1, n.end() + 1) } - for o in edge & nums.keys(): - nums[o].append(int(n.group())) - return nums + + # append each number to the list of numbers + for o in numbers & symbols.keys(): + symbols[o].append(int(n.group())) + return symbols if __name__ == "__main__":