The Cleaning Robot Puzzle: Dirt, Capacity and a Battery

Every floor field gets some dirt, the robot gets a bag that fills up and then a battery that runs out, and the cleaning robot puzzle turns into subset sum, a connected knapsack and the orienteering problem.

In Part 1 the robot could clean everything, so any complete walk was good enough. This part changes the rules so that choices matter: the robot can’t take everything any more, so it has to decide what to take. It builds on Part 1’s depth-first walk and assumes you can read Python.

The new rules

Every floor field now holds some dirt, written as a digit:

The robot now has a capacity: the total amount of dirt its bag can hold. It must never go over capacity, and the goal is to collect as much dirt as possible. Everything else stays as in Part 1: the moves are ^ v < >, rooms are up to 40×40, the walls surround the room, every floor field is reachable, and an answer may have at most 50,000 commands.

There are two pickup rules:

• Take all. Like Part 1, the robot cleans every field it drives onto: stepping onto a dirty field empties all of its dirt into the bag. So the robot may step onto a dirty field only if all of that dirt still fits. Driving back over a field it has already cleaned costs nothing.
• Pick cells. The robot drives with its brushes up, so it can cross dirty fields without collecting anything. A new command, C, cleans the field it is standing on. C is only allowed on a dirty field whose dirt still fits.

And two limits:

• Commands only. At most 50,000 commands, as in Part 1.
• Battery. At most battery movement commands (^ v < >). Cleaning with C uses no battery.

That gives four problems. I wrote five solutions, because the hardest problem gets both a fast version and an exact version, so the two can be compared:

Each file has a solution(room, capacity) or solution(room, capacity, battery) function. Helpers they share live in dirt_common.py, and every file behind the series is listed in the code’s README.

Examples with dirt

The pickup rule makes a difference. Capacity 1:

• Take all: stepping onto the 5 would overflow the bag, and the 1 is behind it. The best answer is to collect nothing: "".
• Pick cells: drive over the 5 with brushes up, then clean the 1: >>C collects 1.

Grabbing whatever fits is a trap. Capacity 6, take all:

A greedy robot takes the 3 on the left, since it fits. Now the 5 no longer fits, the 1 is behind the 5, and the robot is stuck with 3. The best answer goes right instead: >> collects 5 + 1 = 6. Whatever we build has to look further ahead than “take what fits”.

The running example. A 5×7 room with 40 dirt in total:

With capacity 23, the take-all route also drives back over the already-clean 2 (the ^ in the middle) to reach the right-hand side.

Shared building blocks

All five solutions start from the same few helpers in dirt_common.py.

Reading the room. parse turns the strings into the start position and a dictionary from each floor field to its dirt. Walls simply aren’t in the dictionary, so “is this floor?” is just cell in dirt.

MOVES = [(-1, 0, '^', 'v'), (1, 0, 'v', '^'), (0, -1, '<', '>'), (0, 1, '>', '<')]


def parse(room):
    """Return (start, dirt) where dirt maps every floor cell to its dirt (0 for the start)."""
    start, dirt = None, {}
    for r, row in enumerate(room):
        for c, ch in enumerate(row):
            if ch == '#':
                continue
            if ch == '*':
                start = (r, c)
            dirt[(r, c)] = int(ch) if ch.isdigit() else 0
    return start, dirt


def neighbours(cell, allowed):
    """Yield (neighbour, command to go there, command to come back) for neighbours in allowed."""
    r, c = cell
    for dr, dc, go, back in MOVES:
        n = (r + dr, c + dc)
        if n in allowed:
            yield n, go, back

Walking a group of fields. walk is Part 1’s recursion-free depth-first walk, with two changes. It only enters fields in allowed, which can be the whole floor or just a chosen group. It also records order: every field in first-visit order, together with how many moves the robot had made when it first stood there. That index is what lets the solutions insert a C at exactly the right moment, or cut the route after the last useful move.

def walk(start, allowed):
    """Depth-first walk over the cells of `allowed` connected to start, without recursion.

    Returns (moves, order). `order` lists (cell, i) in first-visit order, where i is the
    number of moves made when the robot first stands on the cell.
    """
    seen = {start}
    moves, order, trail = [], [(start, 0)], []
    cell = start
    while True:
        for n, go, back in neighbours(cell, allowed):
            if n not in seen:
                seen.add(n)
                moves.append(go)
                trail.append((back, cell))
                cell = n
                order.append((cell, len(moves)))
                break
        else:                                   # no new neighbour: retrace one step
            if not trail:
                return moves, order
            back, cell = trail.pop()
            moves.append(back)

Turning a walk into an answer. For take all, the route is the walk over the chosen group, cut right after the robot reaches the last new field. For pick cells, the route is the walk with a C added the first time the robot stands on each chosen field, cut right after the last C.

def take_all_route(start, cells):
    """Route that enters every cell of the connected set `cells`, ending at the last new cell."""
    moves, order = walk(start, cells)
    return ''.join(moves[:order[-1][1]])


def route_with_cleaning(moves, order, chosen):
    """Add 'C' when the robot first stands on a chosen cell; drop the moves after the last 'C'."""
    clean_at = {i for cell, i in order if cell in chosen}
    out, end = [], 0
    for i in range(len(moves) + 1):
        if i in clean_at:
            out.append('C')
            end = len(out)
        if i < len(moves):
            out.append(moves[i])
    return ''.join(out[:end])

Pick cells: subset sum in disguise

The observation that makes it easy

With brushes up, the robot can drive to every floor field without collecting anything. So where the dirt lies doesn’t matter at all, only how much each field holds. The problem becomes:

Given the numbers 3, 4, 5, 1, 5, 2, 1, 5, 4, 3, 2, 5, choose some of them with the largest total that is at most the capacity.

That is the classic subset sum problem. In general it’s NP-hard, but it is only hard when the numbers are huge. Here every number is between 1 and 5, and a 40×40 room holds at most 1,444 × 5 = 7,220 dirt, so the totals we care about are small.

Subset sum with one integer

The trick is to store every total we can make as the bits of a single Python integer. Bit s is 1 when some of the numbers seen so far add up to exactly s.

• Start with reach = 1: only bit 0 is set, because choosing nothing makes 0.
• To add a number w, compute reach | (reach << w). The shift moves every old total up by w (“also take this number”), and the OR keeps the old totals (“or don’t”).
• AND with a mask to throw away totals above the capacity.

For example, with the numbers 3, 4, 5 and capacity 8:

The best total is the highest set bit: reach.bit_length() - 1 = 8.

Python integers can be any size, so each step is one fast shift-and-OR over about 7,000 bits. The whole room takes milliseconds.

Finding which numbers made the best total

Keep the integer after every step, then walk backwards from the best total s. At step i, ask: could s already be made without number i? If yes, skip that number. If not, number i must be part of the total, so take it and subtract it from s.

With 3, 4, 5 and s = 8: 8 can’t be made from {3, 4}, so take the 5 (s = 3). 3 can be made from {3}, so skip the 4. 3 can’t be made from nothing, so take the 3 (s = 0). The answer is 3 + 5.

Skipping whenever possible has a nice side effect. The numbers are listed in the order the walk first visits their fields, so the chosen fields end up as early along the walk as possible. That makes the route, which is cut after the last C, as short as possible.

def best_subset(weights, capacity):
    """Choose weights with the largest sum that fits capacity. Returns (sum, chosen indices).

    Bit s of reach[i] is set when some of the first i weights add up to exactly s.
    """
    limit = min(capacity, sum(weights))
    mask = (1 << (limit + 1)) - 1
    reach = [1]
    for w in weights:
        reach.append(reach[-1] | ((reach[-1] << w) & mask))
    best = reach[-1].bit_length() - 1
    chosen, s = [], best
    for i in range(len(weights), 0, -1):
        if not (reach[i - 1] >> s) & 1:         # s can't be made without weight i-1
            chosen.append(i - 1)
            s -= weights[i - 1]
    return best, chosen

One detail matters: the mask is built from min(capacity, total), not from the capacity itself. Nothing stops someone passing a capacity of a billion, and 1 << 1_000_000_001 is a 125 MB integer.

The pick-cells solution

Walk the whole floor, choose the best set of dirty fields, and clean them as the walk passes:

from dirt_common import best_subset, parse, route_with_cleaning, walk


def solution(room, capacity):
    start, dirt = parse(room)
    moves, order = walk(start, dirt)
    dirty = [cell for cell, _ in order if dirt[cell] > 0]
    _, picked = best_subset([dirt[cell] for cell in dirty], capacity)
    return route_with_cleaning(moves, order, {dirty[i] for i in picked})

Why it’s always the best. The robot can reach and clean any combination of fields, and best_subset finds the best combination exactly.

Size and speed. The walk has at most 2,886 moves, plus at most 1,444 Cs. On a 40×40 room it takes about 5 ms.

That settles the easy rule. The other rule is where it gets interesting.

Take all: choosing a connected group

What really changes

Under take all, every field the robot enters is emptied into the bag, and it can’t enter a field whose dirt doesn’t fit. Three things follow:

• The fields it cleans form a connected group that contains the start, because the robot has to drive through the group to reach any part of it.
• The route never leaves that group, because any field outside it would get cleaned too.
• Once the group is chosen, the route is easy: the walk from the shared helpers enters every field of the group in at most 2 × (size − 1) moves.

So the whole problem is: choose a connected group containing the start, with total dirt at most the capacity, and make that total as large as possible.

This is known as the connected knapsack problem, and it’s NP-hard on general graphs. Dirt values of at most 5 help, but I don’t know a fast method that is always exact. So there are two versions: a fast one that is almost always the best, and an exact one to check it against.

Two facts that frame the answer

An upper bound. Any group the take-all robot can collect is also a legal choice for the pick-cells robot. So the pick-cells answer, call it U, is at least as good as the best take-all answer. That makes it a certificate: if a take-all group totals exactly U, nothing can beat it and the search can stop. It catches the common cases for free: everything fits, the bag fills exactly, or the dirt values can’t make a larger total anyway (for example, when every field holds a 2 or a 4, an odd capacity can’t be reached).

Never more than 4 short. Suppose the total dirt is more than the capacity. Grow a group one field at a time, always adding a field next to the group. Each step adds at most 5 dirt, so the totals climb in steps of at most 5. The last group that still fits is therefore within 4 of the capacity. The best answer can’t exceed the capacity, so even a simple method lands within 4 of it.

On a tree, the problem is easy

Suppose for a moment the floor had no loops, like a maze where there is exactly one path between any two fields. Then the floor is a tree hanging from the start, and a connected group containing the start is just the tree with some branches cut off. That can be solved exactly with dynamic programming.

List the fields in preorder: a field, then everything below it, then its next sibling. In that list every field’s subtree forms one unbroken block. Now go through the list, deciding each field in turn:

• Take field i: move on to position i + 1. Its children come next, and they are now allowed.
• Skip field i: none of its descendants can be connected any more, so jump over its whole block to position i + size[i].

dp[i] is the bitset of totals that are possible when we are about to decide field i, using the same bitset trick as the pick-cells solution. Taking field i shifts the bits by its dirt, and skipping it copies the bits forward.

A worked example. A small tree: the start S has two children, A with 3 dirt and B with 5. B has one child, C with 1. Capacity 6. The preorder is S, A, B, C, with block sizes 4, 1, 2, 1.

The best total is 6: S, B and C. To recover the group, walk backwards from the end. At each position, check whether the total came from taking the previous field or from skipping a block that ends here. A small skips_into list records which blocks end where.

def best_subtree(start, parent, children, dirt, capacity):
    """Exact best set that is connected inside this tree, contains start and fits capacity."""
    order, stack = [], [start]
    while stack:                                # preorder: each subtree is one contiguous block
        cell = stack.pop()
        order.append(cell)
        stack.extend(children[cell])
    size = {cell: 1 for cell in order}
    for cell in reversed(order[1:]):
        size[parent[cell]] += size[cell]

    n = len(order)
    w = [dirt[cell] for cell in order]
    span = [size[cell] for cell in order]
    mask = (1 << (capacity + 1)) - 1
    # dp[i]: bit s is set when the choices for order[0..i-1] collect exactly s
    # and order[i] may still be taken (its parent is in the set).
    dp = [0] * (n + 1)
    dp[1] = 1                                   # the start is always in the set, with no dirt
    for i in range(1, n):
        if dp[i]:
            dp[i + 1] |= (dp[i] << w[i]) & mask     # take order[i]
            dp[i + span[i]] |= dp[i]                # skip order[i] and its whole subtree

    best = dp[n].bit_length() - 1
    skips_into = [[] for _ in range(n + 1)]
    for i in range(1, n):
        skips_into[i + span[i]].append(i)
    chosen, i, s = {start}, n, best
    while i > 1:
        if s >= w[i - 1] and (dp[i - 1] >> (s - w[i - 1])) & 1:
            chosen.add(order[i - 1])
            s -= w[i - 1]
            i -= 1
        else:
            i = next(k for k in skips_into[i] if (dp[k] >> s) & 1)
    return best, chosen

It processes about 2 × (number of fields) big integers of at most 7,221 bits, so it runs in a few milliseconds on a 40×40 room.

Real rooms have loops: try many spanning trees

A real room has loops, but here is the key fact: every connected group containing the start is a subtree of some spanning tree of the floor. Take any tree that connects the group, then keep attaching the remaining fields to it. So if we run the tree DP on the right spanning tree, we get the true best. We don’t know which tree is the right one, so we try several.

The trees are grown with a randomized version of Prim’s algorithm. Start from the robot’s field, keep a heap of fields next to the tree, and repeatedly attach the one with the smallest key. Two tricks make the rounds work together.

Trick 1: build the tree around the best group so far. Give each candidate field the key (field is not in the best group, random number). Fields in the best group have False as the first part, so all of them attach before anything else, and the group becomes a subtree of the new tree. The DP on that tree can therefore never do worse than the group we already have. In one pass it tries every way of cutting branches off the group and growing new branches outside it.

Trick 2: sometimes start from scratch. Trick 1 has a blind spot, and testing against brute force found it. Take this room with capacity 20:

The first version of the fast solver returned 16. It had settled on a group whose only link to the bottom row ran through the 1. The best answer, 20, drops the 1 and collects all four 5s by going down the left side instead. But a tree built around the old group always keeps that path through the 1, so cutting branches can never find the new shape. The fix is to make every other round use a completely random tree. If that round finds something better, or equally good, the search moves to it. After that change the fast solver found the best answer on all 800 test rooms that brute force can check, instead of 798.

The rounds stop as soon as the best total reaches the upper bound U, and after at most 30 rounds otherwise. The random generator uses a fixed seed, so the same room always gets the same answer.

import heapq
import random

from dirt_common import best_subset, neighbours, parse, take_all_route


def spanning_tree(start, dirt, keep, rng):
    """Random spanning tree grown from start (Prim). Cells in `keep` are attached before any
    other cell, so a connected `keep` containing start becomes a subtree hanging from the root."""
    parent, children = {start: None}, {start: []}
    heap = []
    cell = start
    while True:
        for n, _, _ in neighbours(cell, dirt):
            if n not in parent:
                heapq.heappush(heap, (n not in keep, rng.random(), n, cell))
        while heap and heap[0][2] in parent:
            heapq.heappop(heap)
        if not heap:
            return parent, children
        _, _, cell, par = heapq.heappop(heap)
        parent[cell] = par
        children[cell] = []
        children[par].append(cell)


def best_connected_set(start, dirt, capacity, rounds=30, seed=0):
    """Return (best, cells, upper).

    cells is a connected set containing start whose dirt, best, fits capacity.
    upper is a bound no connected set can beat, so best == upper proves best is optimal.
    """
    total = sum(dirt.values())
    if total <= capacity:
        return total, set(dirt), total
    upper, _ = best_subset(list(dirt.values()), capacity)
    rng = random.Random(seed)
    best, cells = 0, {start}
    for round_no in range(rounds):
        if best == upper:
            break
        # Even rounds build the tree around the best set, so they can only improve on it.
        # Odd rounds use a fully random tree, which can reshape the set out of a dead end.
        keep = cells if round_no % 2 == 0 else {start}
        parent, children = spanning_tree(start, dirt, keep, rng)
        value, chosen = best_subtree(start, parent, children, dirt, capacity)
        if value >= best:                       # on a tie, move to the new set to keep exploring
            best, cells = value, chosen
    return best, cells, upper


def solution(room, capacity):
    start, dirt = parse(room)
    _, cells, _ = best_connected_set(start, dirt, capacity)
    return take_all_route(start, cells)

What you can rely on.

• Never more than 4 below the best. Even the first round’s tree gives that.
• Exact on mazes. When the floor has no loops, the only spanning tree is the floor itself, so the answer is always the best.
• Proven best whenever it reaches U. It did on 18 of twenty 40×40 test rooms, in 27 ms on average and 286 ms at worst.

Take all, exactly

The exact version starts by running the fast one. If the fast answer reaches the upper bound U, it is already proven best, and we are done. Otherwise it searches through every connected group that could still beat it.

Visiting every connected group exactly once

The search keeps three things:

• the group, which always contains the start;
• the edge, the undecided fields next to the group;
• the banned fields, which may not be added in this part of the search.

It picks one field v from the edge and splits the search in two:

1. Take v (only if its dirt fits). Fields next to v that haven’t been seen yet join the edge.
2. Ban v. For the rest of this branch, v may never be added.

Every connected group is found exactly once. Two different groups disagree about some edge field, and the first such field sends them down different branches.

Cutting the search short

• Hopeless branches. Keep available, the total dirt of every field that is neither in the group nor banned. If total + available can’t beat the best found so far, stop exploring this branch.
• The certificate. As soon as the best reaches U, stop the whole search.
• Fields that can never fit. Any field whose dirt alone is more than the capacity is banned before the search starts.

import sys

from dirt_common import neighbours, parse, take_all_route
from take_all_fast import best_connected_set


def best_connected_set_exact(start, dirt, capacity):
    best, best_cells, upper = best_connected_set(start, dirt, capacity)
    if best == upper:
        return best, best_cells
    sys.setrecursionlimit(max(sys.getrecursionlimit(), 3 * len(dirt) + 100))

    # 'in' = in the set, 'edge' = next to the set and undecided, 'out' = never take; absent = untouched
    status = {cell: 'out' for cell, d in dirt.items() if d > capacity}
    status[start] = 'in'
    taken = [start]
    first_edge = []
    for n, _, _ in neighbours(start, dirt):
        if n not in status:
            status[n] = 'edge'
            first_edge.append(n)
    available = sum(d for cell, d in dirt.items() if status.get(cell) != 'out')

    def search(edge, total, available):
        """Each connected set is reached exactly once: pick an edge cell, then take it or ban it."""
        nonlocal best, best_cells
        if total > best:
            best, best_cells = total, set(taken)
        if best == upper or not edge or total + available <= best:
            return
        v, rest = edge[-1], edge[:-1]
        if total + dirt[v] <= capacity:         # branch 1: take v
            status[v] = 'in'
            taken.append(v)
            added = [n for n, _, _ in neighbours(v, dirt) if n not in status]
            for n in added:
                status[n] = 'edge'
            search(rest + added, total + dirt[v], available - dirt[v])
            for n in added:
                del status[n]
            taken.pop()
        status[v] = 'out'                       # branch 2: never take v
        search(rest, total, available - dirt[v])
        status[v] = 'edge'

    search(first_edge, 0, available)
    return best, best_cells


def solution(room, capacity):
    start, dirt = parse(room)
    _, cells = best_connected_set_exact(start, dirt, capacity)
    return take_all_route(start, cells)

The recursion is at most one level deep per field, since every call decides one field. That’s why the limit is raised to about three times the number of fields. The catch is that the number of connected groups grows exponentially. On a large room where the fast answer can’t be proven best, the exact search could run for a very long time.

Fast vs exact: the comparison

In every room where the exact search could run, the fast version was already the best. It proves that itself in the large majority of rooms, and when it can’t, the exact search has so far only confirmed it. That makes the exact version more of a measuring stick than something you’d run on big rooms.

Adding a battery

With a battery, the length of the route matters too: the robot can make only battery moves. Picking the route that collects the most within a move budget is a form of the orienteering problem, which is NP-hard even without a bag that fills up. So both battery solutions are heuristics: they are usually the best and always legal, but they come with no guarantee.

Both begin with the same shortcut: if the unlimited answer already fits in the battery, return it. It’s the best possible answer, and the battery didn’t get in the way.

Pick cells with a battery

The idea: walk towards whatever is most useful per step, and decide what to clean at the very end.

• Clean afterwards. Cleaning is optional and costs no battery. So once the walk is known, the robot can choose which of the dirty fields it passed to clean. Run best_subset over them in visiting order and insert the Cs. The walk itself only has to visit a good mix of dirt.
• Measure usefulness, not size. Far from full, a 5 is worth 5. Close to full, a 5 may be worth nothing, while a 1 is exactly what’s missing. gain_of[w] is how much one more field with dirt w would raise the best total that could still be collected. It’s computed with the same bitset as before: compare reach | (reach << w) with reach.
• Pick the next target. Run a breadth-first search from the robot. Among the shortest paths to each field, keep the path that passes the most useful uncleaned dirt. Score each field as that useful dirt (capped at the space left in the bag) divided by distance ** alpha. Walk to the field with the best score.
• Try a few temperaments. alpha sets how much distance hurts: 1 is willing to travel for big piles, 2 prefers what’s close. The solution tries 1, 1.5 and 2 and keeps the best result, preferring fewer moves on a tie.
• Stay fast. The search first looks only 6 steps around the robot, and scans the whole room only when nothing nearby helps.

from collections import deque

import pick_cells
from dirt_common import best_subset, neighbours, parse, route_with_cleaning

NEAR = 6        # look this many steps around the robot first; search further only if needed


def moves_used(route):
    return sum(ch != 'C' for ch in route)


def solution(room, capacity, battery):
    route = pick_cells.solution(room, capacity)
    if moves_used(route) <= battery:
        return route                            # the unlimited best already fits
    start, dirt = parse(room)
    best_route, best_value = '', 0
    for alpha in (1, 1.5, 2):
        route, value = greedy(start, dirt, capacity, battery, alpha)
        if value > best_value or (value == best_value and len(route) < len(best_route)):
            best_route, best_value = route, value
    return best_route


def greedy(start, dirt, capacity, battery, alpha):
    """Walk towards the most useful dirt per step, then clean the best subset of what was visited."""
    limit = min(capacity, sum(dirt.values()))
    mask = (1 << (limit + 1)) - 1
    reach = 1                                   # bit s set: some visited cells add up to exactly s
    pos, left = start, battery
    visited = {start}
    moves, order = [], [(start, 0)]

    while left > 0:
        have = reach.bit_length() - 1
        if have == limit:
            break
        # how much one more cell with dirt w would raise the best amount we could collect
        gain_of = [0] + [((reach | (reach << w)) & mask).bit_length() - 1 - have for w in range(1, 6)]
        target, parent = find_target(pos, dirt, visited, gain_of, limit - have, min(left, NEAR), alpha)
        if target is None and left > NEAR:
            target, parent = find_target(pos, dirt, visited, gain_of, limit - have, left, alpha)
        if target is None:
            break
        path, cell = [], target
        while cell != pos:
            prev, go = parent[cell]
            path.append((go, cell))
            cell = prev
        for go, cell in reversed(path):
            moves.append(go)
            left -= 1
            if cell not in visited:
                visited.add(cell)
                order.append((cell, len(moves)))
                reach |= (reach << dirt[cell]) & mask
        pos = target

    dirty = [cell for cell, _ in order if dirt[cell] > 0]
    value, picked = best_subset([dirt[cell] for cell in dirty], capacity)
    return route_with_cleaning(moves, order, {dirty[i] for i in picked}), value


def find_target(pos, dirt, visited, gain_of, still_fits, radius, alpha):
    """BFS up to `radius` steps. Among shortest paths, keep the one passing the most useful
    unvisited dirt, and return the cell with the best useful-dirt-per-step score."""
    dist, gain, parent = {pos: 0}, {pos: 0}, {}
    queue = deque([pos])
    while queue:
        cell = queue.popleft()
        if dist[cell] == radius:
            continue
        for n, go, _ in neighbours(cell, dirt):
            g = gain[cell] + (0 if n in visited else gain_of[dirt[n]])
            if n not in dist:
                dist[n], gain[n], parent[n] = dist[cell] + 1, g, (cell, go)
                queue.append(n)
            elif dist[n] == dist[cell] + 1 and g > gain[n]:
                gain[n], parent[n] = g, (cell, go)

    target, best_score = None, 0
    for cell, g in gain.items():
        if g > 0:
            score = min(g, still_fits) / dist[cell] ** alpha
            if score > best_score:
                target, best_score = cell, score
    return target, parent

One subtle point is in find_target. A field can be reached by several shortest paths, and the search keeps the one with the most useful dirt. Updating a field that was already found is safe, because breadth-first search finishes all fields at distance d before it expands any field at distance d + 1.

Take all with a battery

Under take all, the robot can only drive over fields it has already cleaned, since anything else would be cleaned on the way. So reaching the next field costs battery but no capacity. The solution compares four candidates:

• Candidate 1: the unlimited take-all route, cut off when the battery runs out.
• Candidates 2–4: grow the cleaned area greedily, once each for alpha = 1, 1.5 and 2. Search outward from the robot through cleaned fields. Every uncleaned neighbour whose dirt fits is a candidate, scored as dirt / (steps to reach it) ** alpha. Walk there and clean it. The search stops early once even a 5 that far away couldn’t beat the best score found so far.

from collections import deque

from dirt_common import neighbours, parse, walk
from take_all_fast import best_connected_set


def solution(room, capacity, battery):
    start, dirt = parse(room)
    _, cells, _ = best_connected_set(start, dirt, capacity)
    moves, order = walk(start, cells)
    full = moves[:order[-1][1]]
    if len(full) <= battery:
        return ''.join(full)                    # the unlimited answer already fits

    # candidate 1: follow that route until the battery runs out
    best_route = ''.join(full[:battery])
    best_value = sum(dirt[cell] for cell, i in order if i <= battery)
    for alpha in (1, 1.5, 2):                   # more candidates: grow greedily
        route, value = greedy(start, dirt, capacity, battery, alpha)
        if value > best_value or (value == best_value and len(route) < len(best_route)):
            best_route, best_value = route, value
    return best_route


def greedy(start, dirt, capacity, battery, alpha):
    """Repeatedly step into the uncleaned cell with the most dirt per move that still fits.

    The robot only walks over cells it has already cleaned, so reaching a cell costs
    battery but no capacity.
    """
    cleaned, pos, left, room_left = {start}, start, battery, capacity
    moves = []
    while left > 0:
        target, entry, best_score = None, None, 0
        dist, parent = {pos: 0}, {}
        queue = deque([pos])
        while queue:
            cell = queue.popleft()
            d = dist[cell]
            if d + 1 > left or 5 / (d + 1) ** alpha <= best_score:
                break                           # nothing farther away can score higher
            for n, go, _ in neighbours(cell, dirt):
                if n in cleaned:
                    if n not in dist:
                        dist[n], parent[n] = d + 1, (cell, go)
                        queue.append(n)
                elif dirt[n] <= room_left:
                    score = dirt[n] / (d + 1) ** alpha
                    if score > best_score:
                        target, entry, best_score = n, (cell, go), score
        if target is None:
            break
        parent[target] = entry
        steps, cell = [], target
        while cell != pos:
            cell, go = parent[cell]
            steps.append(go)
        steps.reverse()
        moves.extend(steps)
        left -= len(steps)
        cleaned.add(target)
        room_left -= dirt[target]
        pos = target
    return ''.join(moves), capacity - room_left

The running example, with a battery

Back to the 5×7 room, now with capacity 23. Without a battery, take all needs 8 moves and pick cells needs 9. Here is what each heuristic returns as the battery shrinks, next to the true best from a brute-force search over every possible walk:

A few things stand out:

• Battery 7, pick cells. The unlimited route needs 9 moves, so the heuristic had to find a different one. >CvC>C>C>CvC<C runs right along the middle row, down the right side and back one field: 3 + 5 + 2 + 1 + 5 + 5 + 2 = 23 in exactly 7 moves.
• Battery 5, where both heuristics fall short. Both grab the 4 at the top, a big pile one step away, and then have no useful move left. The best route skips it and keeps going along the middle row: >v>>> collects 3 + 5 + 2 + 1 + 5 = 16.
• Take all is the harder battery problem. At batteries 6 and 7 its greedy trails the best by 4. It can only travel through fields it has already cleaned, so an early detour to the 4 keeps costing moves later on.

Across many rooms

On 800 tiny rooms, with a random battery between 0 and twice the number of floor fields, compared with a brute-force search over every walk:

On 40×40 rooms, take_all_battery took 41 ms on average and 270 ms at worst. pick_cells_battery is slower, since it runs a breadth-first search for every target it walks to: 141 ms on average and about 1 second at worst.

Testing Part 2

All of Part 2 is checked by one script, test_dirt.py. It runs in about a minute.

A simulator that knows every rule

The solutions return strings, and a string is easy to get subtly wrong. So every answer is replayed by a simulator that fails on any broken rule and otherwise returns the dirt collected. The tests then compare numbers, not routes.

STEP = {'^': (-1, 0), 'v': (1, 0), '<': (0, -1), '>': (0, 1)}


def simulate(room, cmds, capacity, battery=None, pick=False):
    """Replay cmds and return the dirt collected; fail on any broken rule."""
    start, dirt = parse(room)
    assert len(cmds) <= 50_000, 'more than 50,000 commands'
    allowed = set(STEP) | ({'C'} if pick else set())
    assert set(cmds) <= allowed, f'unexpected command in {cmds!r}'
    moves = sum(ch in STEP for ch in cmds)
    assert battery is None or moves <= battery, f'{moves} moves, battery {battery}'
    pos, collected, cleaned = start, 0, {start}
    for ch in cmds:
        if ch == 'C':
            assert pos not in cleaned and dirt[pos] > 0, f'C on a clean cell {pos}'
            assert collected + dirt[pos] <= capacity, f'C at {pos} overflows capacity'
            collected += dirt[pos]
            cleaned.add(pos)
            continue
        dr, dc = STEP[ch]
        pos = (pos[0] + dr, pos[1] + dc)
        assert pos in dirt, f'walked into a wall at {pos}'
        if not pick and pos not in cleaned:
            assert collected + dirt[pos] <= capacity, f'entering {pos} overflows capacity'
            collected += dirt[pos]
            cleaned.add(pos)
    return collected

Rooms designed to cause trouble

• Random rooms with 0% to 40% furniture. Floor that can’t be reached is turned into walls.
• Perfect mazes, where the floor has no loops. On these the fast take-all solver must be exact, since the floor is its own spanning tree.
• Dirt mixes: all values 1–5, only 5s, only 2s and 4s, only 3s and 5s, and mostly 5s with a rare 1. The restricted mixes leave totals that can’t be made, which is where a solver can get stuck just short of the capacity.
• Capacities: 1, a random value, total − 1, exactly the total, and more than the total.
• Batteries: from 0 up to twice the number of floor fields.

Brute force as the referee

On tiny rooms, with at most 12 floor fields, three independent brute-force searches compute the true best:

• Take all: list every connected group containing the start, one field at a time.
• Pick cells: try every count of 1s, 2s, 3s, 4s and 5s.
• Battery: a breadth-first search over (robot position, set of fields visited) for up to battery moves.

Each search could have bugs of its own, so they are checked against each other. With a battery of four times the number of fields, the walk search must agree with the other two.

Fixed rooms for known traps

• The greedy-trap corridor #3*51# with capacity 6 must collect 6.
• A room of 5s with capacity 12 must collect 10.
• A room where nothing fits must return "".
• A room that is only the start must return "".
• A battery of 0 must return "".

The results

What Part 2 teaches

• Small numbers are a gift. Dirt of at most 5 means totals of at most 7,220, so “every total we could make” fits in a single Python integer, and subset sum becomes a shift and an OR.
• Find an upper bound you can actually reach. The pick-cells answer caps the take-all answer. Reaching the cap turns a heuristic’s answer into a proof, and it lets the exact search stop early.
• Hard on graphs, easy on trees. Connected knapsack is hard in general, but on a tree it’s a clean dynamic program. Random spanning trees carry that exact solver over to rooms with loops.
• Test against brute force, not just examples. The first fast solver passed every hand-made room. Only a comparison with brute force on hundreds of tiny rooms turned up the room where it stopped at 16 instead of 20, and that room led directly to the fix.
• Be honest about heuristics. With a battery, the solutions are usually the best but not always, and the running example shows exactly where a greedy step goes wrong.