Advent of Code 2024, day 04

AdventOfCode2024.Day04.py

import re

with open("./data/input04.txt") as inputFile:
    input = inputFile.read().splitlines()

upper_bound = len(input)-1

def rev(ar): return ar[::-1]

# Part 1

def count_xmas(line):
    return len(re.findall(r'XMAS', line)) + len(re.findall(r'XMAS', rev(line)))

res = 0

# Horizontal
for line in input:
    res += count_xmas(line)

# Vertical
for line in [list(row) for row in zip(*rev(input))]:
    line = ''.join(line)
    res += count_xmas(line)

# Diagonal
for i in range(upper_bound+1):
    line1 = ''
    line2 = ''
    line3 = ''
    line4 = ''
    for j in range(0, i+1): 
        line1 += input[j][i-j]
        line2 += input[j][upper_bound-i+j]
        if i < upper_bound:
            line3 += input[upper_bound-j][upper_bound-i+j]
            line4 += input[upper_bound-j][i-j]
    res += count_xmas(line1)
    res += count_xmas(line2)
    res += count_xmas(line3)
    res += count_xmas(line4)

print(res)

# Part 2

def is_xmas(i, j):
    if input[i][j] == 'A' and i > 0 and j > 0 and i < upper_bound and j < upper_bound:
        return ((input[i-1][j-1] == 'M' and input[i+1][j+1] == 'S' or input[i-1][j-1] == 'S' and input[i+1][j+1] == 'M') and 
                (input[i-1][j+1] == 'M' and input[i+1][j-1] == 'S' or input[i-1][j+1] == 'S' and input[i+1][j-1] == 'M'))
    else:
        return False

res = 0

for i in range(upper_bound+1):
    for j in range(upper_bound+1):
        if is_xmas(i, j): 
            res += 1

print(res)