"""Pick cells: the robot may drive over dirt and cleans a cell only with 'C'. Exact."""
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})