day25/solution_jerpint.py

def load_data(file):
    with open(file) as f:
        data = f.readlines()

    # Add newline to make it easier to parse
    data = data + ["\n"]

    keys = []
    locks = []


    grid = []
    for idx, line in enumerate(data):
        line = line.strip("\n")

        if len(line) == 0:
            if grid[0][0] == "#":
                locks.append(grid)
            else:
                keys.append(grid)
            grid = []
            continue

        grid.append(line)

    return keys, locks



def get_lock_profile(lock):
    M = len(lock)
    N = len(lock[0])
    profile = []
    for j in range(N):
        for i in range(M):
            if lock[i][j] == ".":
                profile.append(i-1)
                break
    return profile


def get_key_profile(key):
    M = len(key)
    N = len(key[0])
    profile = []
    for j in range(N):
    #  for j in range(N-1, -1, -1):
        for i in range(M):
            if key[i][j] == "#":
                profile.append(6-i)
                break
    return profile


def check_fit(key, lock):
    for k, l in zip(key, lock):
        if k + l > 5:
            return False
    return True


#  file = "test.txt"
file = "input.txt"
keys, locks = load_data(file)

key_profiles = []
for key in keys:
    key_profiles.append(get_key_profile(key))

lock_profiles = []
for lock in locks:
    lock_profiles.append(get_lock_profile(lock))

fits = []
for key in key_profiles:
    for lock in lock_profiles:
        fits.append(check_fit(key, lock))


print(sum(fits))