adventofcode/2024/22/main.py

#!/usr/bin/env python
# -*- encoding: utf-8; py-indent-offset: 4 -*-

# Author: Sebastian Mark
# CC-BY-SA (https://creativecommons.org/licenses/by-sa/4.0/deed.de)

# pylint: disable=missing-module-docstring,missing-function-docstring,consider-using-f-string


def readinput():
    with open("input", "r", encoding="utf-8") as file:
        lines = [int(line.strip()) for line in file.readlines()]
    return lines


def mix_and_prune(i: int, j: int) -> int:
    return (i ^ j) % 16777216


def calc_secret(num: int) -> (int, list, list):
    changes = []
    numbers = []
    prev_num = num

    for _ in range(2000):
        # step 1
        tmp = num * 64
        num = mix_and_prune(num, tmp)

        # step 2
        tmp = num // 32
        num = mix_and_prune(num, tmp)

        # step 3
        tmp = num * 2048
        num = mix_and_prune(num, tmp)

        numbers.append(num)
        changes.append(num % 10 - prev_num % 10)
        prev_num = num

    return num, numbers, changes


def find_in_list(lst: list, seq: list) -> int:
    seq_len = 4
    for i in range(len(lst) - seq_len + 1):
        if lst[i : i + seq_len] == seq:
            return i
    return -1


def get_subsequences(lsts):
    sequences = set()
    length = 4
    for lst in lsts:
        sequences.update(
            tuple(lst[i : i + length]) for i in range(len(lst) - length + 1)
        )
    return sequences


def main():
    intial_secrets = readinput()

    # part 1
    new_secrets = {}
    for num in intial_secrets:
        new_secrets[num], _, _ = calc_secret(num)
    print("Sum: %d" % sum(new_secrets.values()))

    # part 2 (very ugly and very slow brute force)
    secret_list = {}
    changes = {}
    for num in intial_secrets:
        _, secret_list[num], changes[num] = calc_secret(num)

    bananas = {}
    sequences = get_subsequences(changes.values())
    for i, sequence in enumerate(sequences):
        print(f"\rProcessing sequence {i}/{len(sequences) - 1}", end="", flush=True)
        for secret, change_values in changes.items():
            seq_pos = find_in_list(change_values, list(sequence))
            if seq_pos != -1:
                bananas[sequence] = (
                    bananas.get(sequence, 0) + secret_list[secret][seq_pos + 3] % 10
                )
    print("\nBananas: %d " % max(bananas.values()))


if __name__ == "__main__":
    main()