rooks.py

# Pos is row-major unfolded grid
# To speed up, cut off if above best so far.
def search():
    start = [1, 1, 0, 1, 0, 2, 0, 2, 2]
    current_dist = [start]
    next_dist = []
    dist = 0
    seen = set([tuple(start)])

    while current_dist:
        for pos in current_dist:
            if pos == [2, 2, 0, 2, 0, 1, 0, 1, 1]:
                return dist
            next_poses = []
            for i in range(9):
                if pos[i] != 0:
                    # Up
                    next_i = i
                    while next_i >= 3:
                        next_i -= 3
                        if pos[next_i] == 0:
                            next_pos = pos[:]
                            next_pos[next_i] = pos[i]
                            next_pos[i] = 0
                            next_poses.append(next_pos)
                        else:
                            break
                    # Down
                    next_i = i
                    while next_i < 6:
                        next_i += 3
                        if pos[next_i] == 0:
                            next_pos = pos[:]
                            next_pos[next_i] = pos[i]
                            next_pos[i] = 0
                            next_poses.append(next_pos)
                        else:
                            break
                    # Left
                    next_i = i
                    while next_i % 3 != 0:
                        next_i -= 1
                        if pos[next_i] == 0:
                            next_pos = pos[:]
                            next_pos[next_i] = pos[i]
                            next_pos[i] = 0
                            next_poses.append(next_pos)
                        else:
                            break
                    # Right
                    next_i = i
                    while next_i % 3 != 2:
                        next_i += 1
                        if pos[next_i] == 0:
                            next_pos = pos[:]
                            next_pos[next_i] = pos[i]
                            next_pos[i] = 0
                            next_poses.append(next_pos)
                        else:
                            break
            for next_pos in next_poses:
                if tuple(next_pos) not in seen:
                    seen.add(tuple(next_pos))
                    next_dist.append(next_pos)
        dist += 1
        current_dist = next_dist
        next_dist = []
    return None

print(search())