Advent of Code 2023 - Day 1

from string import digits

NUM_STR_MAP = {
    "one": "1",
    "two": "2",
    "three": "3",
    "four": "4",
    "five": "5",
    "six": "6",
    "seven": "7",
    "eight": "8",
    "nine": "9",
}

STR_INT_MAP = {str(i): i for i in range(1, 10)}

NUM_CHARS = set(digits)
NUM_STRS = set(NUM_STR_MAP.keys())

input_data = """..."""

part_a_test_data = """1abc2
pqr3stu8vwx
a1b2c3d4e5f
treb7uchet"""


def part_a(data):
    total = 0
    for row in data.split("\n"):
        nums = [c for c in row if c in NUM_CHARS]
        total += int(nums[0] + nums[-1])
    return total


def make_trie(words):
    root = dict()
    for word in words:
        current_dict = root
        for letter in word:
            current_dict = current_dict.setdefault(letter, {})
    return root


part_b_test_data = """two1nine
eightwothree
abcone2threexyz
xtwone3four
4nineeightseven2
zoneight234
7pqrstsixteen"""


def part_b(data):
    total = 0
    numbers = [*NUM_STR_MAP.values(), *NUM_STR_MAP.keys()]
    trie = make_trie(numbers)
    for word in data.splitlines():
        word_len = len(word)
        word_numbers = []
        for i in range(len(word)):
            j = i
            current_str = ""
            current = trie
            if word[j] not in current:
                i += 1
                continue
            while j < word_len:
                if word[j] not in current:
                    i += 1
                    break
                current_str += word[j]
                if not current[word[j]]:
                    if len(current_str) == 1:
                        word_numbers.append(current_str)
                    else:
                        word_numbers.append(NUM_STR_MAP[current_str])
                    i += 1
                    break
                else:
                    current = current[word[j]]
                    j += 1
        number_str = word_numbers[0] + word_numbers[-1]
        total += int(number_str)
    return total


if __name__ == "__main__":
    print(part_a(data=part_a_test_data))  # 142
    print(part_a(data=input_data))
    print(part_b(part_b_test_data))  # 281
    print(part_b(data=input_data))