adventofcode/2023/day3/main.py

from functools import reduce


def part1(file: str):
    acc = []
    with open(file) as f:
        for line in f:
            acc.append([(c, False) for c in line if c != "\n"])

    for i1, r in enumerate(acc):
        for i2, c in enumerate(r):
            if not c[0].isdigit():
                continue

            check = False

            def checker(a: int, b: int, check: bool):
                # If already true
                if check:
                    return True

                # Sanity check
                if a < 0 or len(acc) <= a:
                    return check
                if b < 0 or len(acc[a]) <= b:
                    return check

                # Adjacency check
                it = acc[a][b][0]
                if not it.isdigit() and not it == ".":
                    return True
                return False

            # Coordinates
            for i in range(-1, 2):
                for j in range(-1, 2):
                    if i == 0 and j == 0:
                        continue
                    check = checker(i1 + i, i2 + j, check)

            if check:
                acc[i1][i2] = (c[0], True)

    # Find problems
    buffer = ""
    res = 0
    for lines in acc:
        test = False
        for i, (c, t) in enumerate(lines):
            if c.isdigit():
                buffer += c
                if t:
                    test = True

            # End of number
            if not c.isdigit() or i + 1 >= len(lines):
                if len(buffer) > 0:
                    if test:
                        res += int(buffer)
                test = False
                buffer = ""

    return res


def part2(file: str):
    acc = []
    default = (-1, 0)
    bad = (-2, 0)
    with open(file) as f:
        for line in f:
            acc.append([(c, default) for c in line if c != "\n"])

    for i1, r in enumerate(acc):
        for i2, c in enumerate(r):
            if not c[0].isdigit():
                continue

            def checker(a: int, b: int, check: tuple[int, int]):
                # Sanity check
                if a < 0 or len(acc) <= a:
                    return check
                if b < 0 or len(acc[a]) <= b:
                    return check

                # Adjacency check
                if acc[a][b][0] == "*":
                    # If already set
                    if check[0] >= 0 or check[0] == -2:
                        # Shouldn't be multiplied more than 2, if already encountered:
                        #   Ignore the number
                        return bad
                    return a, b

                return check

            # Coordinates
            check = default
            for i in range(-1, 2):
                for j in range(-1, 2):
                    if i == 0 and j == 0:
                        continue
                    check = checker(i1 + i, i2 + j, check)

            acc[i1][i2] = (c[0], check)

    # Find problems
    buffer = ""
    res = []
    for lines in acc:
        test = default
        for i, (c, star_pos) in enumerate(lines):
            if c.isdigit():
                buffer += c
                # test déjà défini + différente star
                if test != default and star_pos != test and star_pos != default:
                    test = bad
                # test pas défini + star ok
                if test == default and star_pos != bad:
                    test = star_pos

            # End of number
            if not c.isdigit() or i + 1 >= len(lines):
                if len(buffer) > 0:
                    if test != default and test != bad:
                        res.append((int(buffer), test))
                test = default
                buffer = ""

    ratios = {}
    for number, values in res:
        try:
            ratios[f"{values}"].append(number)
        except KeyError:
            ratios[f"{values}"] = [number]
    return sum(
        reduce(lambda x, y: x * y, gear) for gear in ratios.values() if len(gear) > 1
    )


if __name__ == "__main__":
    assert part1("example.txt") == 4361
    print(part1("input.txt"))

    assert part2("example.txt") == 467835
    print(part2("input.txt"))