aoc_2024/four/four.py

111 lines
2.4 KiB
Python
Raw Permalink Normal View History

2024-12-08 17:52:54 +00:00
import pprint
import pandas as pd
import numpy as np
mas_patterns = [
pd.DataFrame(
[
["M", ".", "M"],
[".", "A", "."],
["S", ".", "S"],
]
),
pd.DataFrame(
[
["M", ".", "S"],
[".", "A", "."],
["M", ".", "S"],
]
),
pd.DataFrame(
[
["S", ".", "M"],
[".", "A", "."],
["S", ".", "M"],
]
),
pd.DataFrame(
[
["S", ".", "S"],
[".", "A", "."],
["M", ".", "M"],
]
),
]
def get_diagonals(input_df: pd.DataFrame):
n_cols = len(input_df.columns)
n_rows = len(input_df)
diagonals = []
for i in range(-1 * n_rows + 1, n_cols):
diagonals.append("".join(np.diag(input_df, i)))
return diagonals
def get_sub_arrays(input_df: pd.DataFrame):
n_cols = len(input_df.columns)
n_rows = len(input_df)
sub_arrays = []
for i in range(n_rows - 2):
for j in range(n_cols - 2):
sub = input_df.iloc[i : i + 3, j : j + 3].copy()
sub.iloc[0, 1] = "."
sub.iloc[1, 0] = "."
sub.iloc[1, 2] = "."
sub.iloc[2, 1] = "."
sub_arrays.append(sub)
return sub_arrays
with open("./four.txt") as f:
wordsearch_input = f.readlines()
lines = [list(s.strip()) for s in wordsearch_input if s]
lines = pd.DataFrame(lines)
possible_strings = []
# Rows
possible_strings.extend("".join(line.to_list()) for _, line in lines.iterrows())
possible_strings.extend("".join(line.to_list()) for _, line in lines.items())
# Diagonals tl - br
possible_strings.extend(get_diagonals(lines))
# Diagonals bl - tr
possible_strings.extend(get_diagonals(pd.DataFrame(np.rot90(lines))))
# Find xmas
occurences = 0
for string in possible_strings:
occurences += string.upper().count("XMAS")
occurences += string.upper().count("SAMX")
print(occurences)
sub_arrays = get_sub_arrays(lines)
better_occurences = 0
for i, sub_array in enumerate(sub_arrays):
for j, pattern in enumerate(mas_patterns):
try:
equality = (sub_array.values == pattern.values).all()
if equality:
better_occurences += 1
break
except Exception:
print(i, j)
print(sub_array)
print(pattern)
raise
print(better_occurences)