advent-of-code/2024/python/day04.py

import matrix
import shared
from collections import Counter

# XMAS = ['X','M','A','S']
XMAS = "XMAS"


# @shared.profile
def part1(mx):
    count = 0
    for idr, row in enumerate(mx):
        for idc, c in enumerate(row):
            los = matrix.line_of_sight(mx, idr, idc, distance=4, diagonals=True)
            words = list(map("".join, los.values()))
            counter = Counter(words)
            count += counter[XMAS]
    print(count)


# @shared.profile
def part2(mx):
    count = 0
    for idr, row in enumerate(mx):
        for idc, c in enumerate(row):
            if c != "A":
                continue

            nb = matrix.get_neighbors_cardinal(mx, idc, idr, diagonals=True)

            eastdown= []
            for d in ("NW","SE"):
                eastdown.append(nb.get(d, {}).get("value"))
            if None in eastdown:
                continue
            eastdown = sorted(eastdown)

            westdown = []
            for d in ("NE", "SW"):
                westdown.append(nb.get(d, {}).get("value"))
            if None in westdown:
                continue
            westdown = sorted(westdown)

            if eastdown == ['M','S'] and westdown == ['M', 'S']:
                count += 1


    print(count)

def main():
    mx = matrix.load_matrix_file(shared.get_fname(4), str)
    with shared.elapsed_timer() as elapsed:
        part1(mx)
        print("🕒", elapsed())

    mx = matrix.load_matrix_file(shared.get_fname(4), str)
    with shared.elapsed_timer() as elapsed:
        part2(mx)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()