def parse_schematic(schematic_lines): heights = [] for col in range(len(schematic_lines[0])): height = 0 for row in range(len(schematic_lines)): if schematic_lines[row][col] == '#': break height += 1 heights.append(height) return heights def check_fit(lock_heights, key_heights): for i in range(len(lock_heights)): if lock_heights[i] + key_heights[i] > 5: return False return True file = "input.txt" with open(file) as f: lines = f.read().splitlines() locks = [] keys = [] current_schematic = [] for line in lines: if line.startswith("#") or line.startswith("."): current_schematic.append(line) elif line == "": if current_schematic[0].startswith("#"): locks.append(parse_schematic(current_schematic)) else: keys.append(parse_schematic(current_schematic)) current_schematic = [] if current_schematic: # Handle last schematic if current_schematic[0].startswith("#"): locks.append(parse_schematic(current_schematic)) else: keys.append(parse_schematic(current_schematic)) fitting_pairs = 0 for lock in locks: for key in keys: if check_fit(lock, key): fitting_pairs += 1 print(fitting_pairs) print("There is no part 2 for this challenge.")