AoC 2025, Task 10

task10.py

from functools import cmp_to_key
from typing import Optional

# `inp` should be set to actual task input (a multiline string)
lines = [line.split(" ") for line in inp.split("\n")]
test_cases = []
for line in lines:
    test_cases.append({
        "goal": [let == "#" for let in line[0][1:-1]],
        "buttons": [[int(val) for val in prnt[1:-1].split(",")] for prnt in line[1:-1]],
        "jolt": tuple([int(val) for val in line[-1][1:-1].split(",")])
    })
test_cases

Button = list[int]

calls = 0
skipped = 0
use_seen_states = False  # Not using this optimisation since the gains are minimal, see the comment at the bottom
def knap(
    state: tuple[int],  # index |-> current count
    jolt: tuple[int],  # index |-> target count
    btns: list[Button],  # Sorted list of buttons (see the main function explaining the sort), each button being list of indices that are incremented by clicking the button
    presses: int,  # How many button presses in total so far 
    cur_best: Optional[int],  # How many presses min so far (we are doing a BFS search, each successful branch being one set of button presses that leads to jolt)
    seen_states: dict[tuple[int], int],  # What's the current min button press count for each seen state?
):
    global calls, skipped
    calls += 1

    # Two checks/optimisations in the two loops below:
    # 1. Is there an obvious button to max-out, i.e. state and jolt differ at i and this button is the only one affecting i? If so, max out this button and remove it from rec. calls.
    # 2. Is there an index i such that jolt[i] - state[i] is more buttons than what we would need to be better than the current best solution? (early exit if so)
    # There is quite a bit of repetition between this block of code and the remainder of the function; the remainder of the function has better comments.
    freq = [0] * 10
    last_good_button = [0] * 10
    for btn_index, btn in enumerate(btns):
        for b in btn:
            freq[b] += 1
            last_good_button[b] = btn_index
    for i in range(10):
        if cur_best and i < len(jolt) and jolt[i] - state[i] > 0 and presses + (jolt[i] - state[i]) >= cur_best:
            return cur_best
            
        if freq[i] == 1 and state[i] < jolt[i]:
            inc_step = jolt[i] - state[i]
            cstate = list(state)
            btn_index = last_good_button[i]
            btn = btns[btn_index]
            for b in btn:
                cstate[b] += inc_step
                if cstate[b] > jolt[b]:
                    return cur_best
            t_state = tuple(cstate)
            t_presses = presses + inc_step
            
            if use_seen_states and t_state in seen_states and seen_states[t_state] < t_presses:
                skipped += 1
                res = cur_best
            else:
                if t_state == jolt:
                    return t_presses
                elif cur_best is not None and t_presses >= cur_best - 1:
                    return cur_best
                else:
                    res = knap(t_state, jolt, btns[:btn_index] + btns[btn_index + 1:], t_presses, cur_best, seen_states)
                    if not cur_best or (res is not None and res < cur_best):
                        return res
                    return cur_best
        
    # Where do state and jolt differ? In this call we only try to make them closer at first such index
    diff_index = 0
    while True:
        if state[diff_index] > jolt[diff_index]:
            # invalid state
            return cur_best
        if state[diff_index] < jolt[diff_index]:
            break
        diff_index += 1

    best = cur_best
    for btn_index, btn in enumerate(btns):
        cstate = list(state)

        # As mentioned above, we only care about buttons whose first affected index affects the first state/jolt diff
        if btn[0] > diff_index:
            break

        bad_btn = False
        inc_step = 1
        
        # We can skip some recursive calls if this is the last button touching this index; this optimisation is likely not needed anymore
        if btn_index == len(btns) - 1 or btn[0] not in btns[btn_index + 1]:
            inc_step = jolt[diff_index] - state[diff_index]
        for b in btn:
            cstate[b] += inc_step
            if cstate[b] > jolt[b]:
                bad_btn = True
                break
        if bad_btn:
            continue
        t_state = tuple(cstate)
        t_presses = presses + inc_step
        
        if use_seen_states and t_state in seen_states and seen_states[t_state] < t_presses:
            skipped += 1
            res = best
        else:
            if use_seen_states:
                seen_states[t_state] = t_presses
            if t_state == jolt:
                res = t_presses
            elif best is not None and t_presses >= best - 1:
                res = best
            else:
                res = knap(t_state, jolt, btns[btn_index:], t_presses, best, seen_states)
        if not best or (res is not None and res < best):
            best = res

    return best

sm = 0  # The task answer, i.e., the sum of fewest possible button presses in each task
total_calls = 0  # Tracking how many recursive calls in total being made
for case_ind, case in enumerate(test_cases):
    calls = 0
    print("task", case_ind)
    jolt = case["jolt"]
    btns_raw = case["buttons"]

    freq = {}  # index |-> how many buttons increment state at this index
    for i in range(len(jolt)):
        f = 0
        for btn in btns_raw:
            if i in btn:
                f += 1
        freq[i] = f
    
    # We will rearrange indices (reorder jolt and rename indices contained in buttons) so that the new index 0 refers to the slot that least frequently appears in buttons,
    # the idea being we'd like to minimise the number of iterations (possible initial buttons) in the outermost 'loops' of our BFS
    # (the BFS will first do all the button presses for the index 0, then all the remaining button presses involving index 1, etc.)

    # The list containing (original_index, jolt at original index) sorted as explained above
    jolt_w_ind = sorted(list(enumerate(jolt)), key=lambda pair: (freq[pair[0]], pair[1]))

    # Keeping track of new indices so we can rename the indices within button objects
    old_index_to_new_index = {}
    for i, j in enumerate(jolt_w_ind):
        old_index_to_new_index[j[0]] = i
    print("index map", old_index_to_new_index)
    
    # jolt, but reordered using the new indexes, e.g., [1, 2, 3] might become [2, 1, 3] if 1 is the rarest index across buttons (so jolt_renamed[0] := jolt[1] = 2)
    jolt_renamed = tuple([j[1] for j in jolt_w_ind])
    print("jolt", jolt_renamed)

    # As mentioned above, we first want to press all the required buttons that mention 0, after that, value at 0 should match jolt.
    # So, here we sort buttons in a way that ensures all the relevant buttons (containing the index we are optimising for at the moment) are near the beginning of the list.
    # This might make the button selection process a bit easier. Additionally, we renamed the indices; e.g., if 1 is the rarest index across buttons, all buttons containing 1 will now contain 0
    def cmp(btna, btnb):
        for jolt_ind, jolt in jolt_w_ind:
            if jolt_ind in btna and jolt_ind not in btnb:
                return -1
            if jolt_ind in btnb and jolt_ind not in btna:
                return 1
        return 0
    btns = sorted(btns_raw, key=cmp_to_key(cmp)) # lambda btn: 1000 - len(btn))
    for btn_ind in range(len(btns)):
        btns[btn_ind] = sorted([old_index_to_new_index[val] for val in btns[btn_ind]])
    print("buttons", btns)

    # Now we can solve the equivalent task but with a nicer outer-to-inner loop setup
    best = knap((0,) * len(jolt), jolt_renamed, btns, 0, None, {})
    sm += best
    total_calls += calls
    print("best for this case:", best, "total so far:", sm, "calls for this case:", calls, "total calls:", total_calls)
    
print(sm, total_calls)

"""
seen_states optimisation: it turns out that even though this cuts down the number of calls significantly (4x+ on small cases, likely exponential on larger cases), 
the cost of using this data structure is very high, and the gain is not that much, at least in the simple test cases.

However, while I still believed this data structure to be useful, I wrote a proof that this optimisation is sound, i.e., we don't prune any branch we should not be pruning.

Claim:
Let x_i and y_i be buttons and x_i <= x_{i + 1} for all i w.r.t. our modified sort, similarly for y_i. So, each x_i and y_i can be thought of a single button press.
Suppose x_1 + ... + x_n = y_1 + ... + y_{n + k} with k > 1. 
We claim we can discard any solution starting with y_1, ..., y_{n + k}. 

Note: This might not be obvious because what if e.g. x_n is a 'bad' button which can't be built around to produce a good solution, 
while the y_i sequence, though longer, might be extendable to a good solution.

Proof:
It suffices to prove the case k = 1, since other cases follow by induction. 
Suppose the y_i is a part of some sequence of button presses y_1, ..., y_{n + m} for some m (m > 0) such that 
y_1 + ... + y_{n + m} is the best solution (i.e., y_1, ..., y_{n + 1} should not be discarded, contrary to the claim). 
Denote by z_1, ..., z_{m - 1} the possibly empty postfix of this extended sequence which represents 
the button presses that produce the correct solution starting from the earlier sequence of y_i. 
Since x_1 + ... + x_n = y_1 + ... + y_{n + k}, the value of x_1 + ... + x_n + z_1 + ... + z_{m - 1} must be a good solution as well. 
However, the values in this last sum may not be ordered (it is not ordered if and only if x_n > z_1) so it might not be a valid solution. 
So, let w_i := sorted(x_1, ..., x_n, z_1, ..., z_{m - 1}). 
Now, the sum of w_i produces the right number, the w_i as a sequence is a valid sequence of button presses, 
so the assumption that y_i is a part of the best solution is false given it's longer than the w_i sequence by a single button press. 
This proves the claim. 
Note this does not prove that x_1 + ... + x_n is extendable to any correct solution, it might indeed not be, 
the claim only shows we can discard longer solutions once we find shorter ones.
"""