advent/2024/4/one.py

96 lines
2.4 KiB
Python

from itertools import chain
import re
filename = 'sample'
input = [list(i.rstrip()) for i in open(filename).readlines()]
def generate_horizontal_cross_sections(search):
return search
def generate_vertical_cross_sections(search):
output = []
for row in range(len(input[0])):
current = []
for line in search:
current.append(line[row])
output.append(current)
return output
def generate_diagonal_cross_sections(search):
modulo = len(search[0])
concat = ''.join(chain(*search))
output = []
for offset in range(len(concat)):
current = []
index = offset
while index % modulo != 0 and index < len(concat):
current.append(concat[index])
index += modulo - 1
try:
current.append(concat[index])
except IndexError:
pass
output.append(current)
output = [i for i in output if len(i) >= 4]
return output
def generate_antidiagonal_cross_sections(search):
# start from the points on the right and bottom edges
# of the array, and then work northwest diagonally
indexes_to_start = []
modulo = len(search[0])
concat = ''.join(chain(*search))
index = modulo - 1
while index < len(concat):
indexes_to_start.append(index)
index += modulo
index -= modulo
while index % modulo != 0:
indexes_to_start.append(index)
index -= 1
indexes_to_start.append(index)
# work up northwest
output = []
for index in indexes_to_start:
current = []
while index % modulo != 0 and index > 0:
current.append(concat[index])
index -= modulo + 1
if index > 0: current.append(concat[index])
output.append(current)
output = [i for i in output if len(i) >= 4]
return output
search_space = [
*generate_horizontal_cross_sections(input),
*generate_vertical_cross_sections(input),
*generate_diagonal_cross_sections(input),
*generate_antidiagonal_cross_sections(input),
]
xmas = re.compile('XMAS')
total = 0
for item in search_space:
joined = ''.join(item)
print(joined)
results = xmas.findall(joined)
print(results)
total += len(results)
joined = ''.join(reversed(item))
print(joined)
results = xmas.findall(joined)
print(results)
total += len(results)
print(results)