Solver for The Watness 2 (re 450) from PlaidCTF 2020

The Watness 2 solver.py

def split_by(data, cnt):
    return [data[i : i+cnt] for i in range(0, len(data), cnt)]

MAX_PATH = 150  # just some guessed estimate, should be fine

LEVELS = [
    'rbrr rgb rb  r brgrbrgb  grrgbbg grg bgrg  bbgrbg',
    'rrbrb rg g  bgrbgggr ggrgr gr rg brr  b  bggrbgbb',
    'rbr  bbggrgrggb   bggbb b  b bbrbbgg gbrrbgrbbb g',
]

def parse_map(letters):
    return [list(x) for x in split_by(letters, 7)]

def print_map(m):
    for line in m:
        print(''.join(line))

def count_nbh(m, x, y, color):
    res = 0
    #for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
    for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1), (-1, -1), (1, -1), (-1, 1), (1, 1)]:
        tmp_x = x + dx
        tmp_y = y + dy
        if tmp_x < 0 or tmp_y < 0 or tmp_x >= 7 or tmp_y >= 7:
            continue
        if m[tmp_y][tmp_x] == color:
            res += 1
    return res
        
def step_map(m):
    new_map = [[None] * 7 for _ in range(7)]
    for iy in range(7):
        for ix in range(7):
            #empty_cnt = count_nbh(m, ix, iy, ' ')
            red = count_nbh(m, ix, iy, 'r')
            green = count_nbh(m, ix, iy, 'g')
            blue = count_nbh(m, ix, iy, 'b')
            if m[iy][ix] == ' ':
                if green == 0 and blue == 0:
                    new = ' '
                elif blue < green:
                    new = 'g'
                else:
                    new = 'b'
            elif m[iy][ix] == 'r':
                if red not in [2,3]:
                    new = ' '
                elif blue == 0 or green == 0:
                    new = ' '
                else:
                    new = 'r'
            elif m[iy][ix] == 'g':
                if red > 4:
                    new = ' '
                elif blue > 4:
                    new = 'b'
                elif red in [2,3]:
                    new = 'r'
                else:
                    new = 'g'
            elif m[iy][ix] == 'b':
                if red > 4:
                    new = ' '
                elif green > 4:
                    new = 'g'
                elif red in [2,3]:
                    new = 'r'
                else:
                    new = 'b'

            new_map[iy][ix] = new
    return new_map

def is_red(m, x, y):
    if x < 0 or y < 0 or x >= 7 or y >= 7:
        return False
    return m[y][x] == 'r'

def dfs(steps, vis, cur_depth, x, y):
    if (x, y) == (7, 7):
        return ''
    assert vis[y][x] == False
    vis[y][x] = True
    for dx, dy, ch in [(-1, 0, 'L'), (1, 0, 'R'), (0, -1, 'U'), (0, 1, 'D')]:
        new_x = x + dx
        new_y = y + dy
        if new_x < 0 or new_y < 0 or new_x > 7 or new_y > 7:
            continue
        min_x = min(x, new_x)
        min_y = min(y, new_y)
        ok = False
        if new_x != x:
            if is_red(steps[cur_depth], min_x, y - 1) or is_red(steps[cur_depth], min_x, y):
                ok = True
        else:
            assert new_y != y
            if is_red(steps[cur_depth], x - 1, min_y) or is_red(steps[cur_depth], x, min_y):
                ok = True
        if ok and not vis[new_y][new_x]:
            res = dfs(steps, vis, cur_depth + 1, new_x, new_y)
            if res is not None:
                return ch + res
    vis[y][x] = False
    return None

def solve(level):
    m = parse_map(level)
    steps = [m]
    print_map(m)
    for _ in range(MAX_PATH):
        m = step_map(m)
        steps.append(m)
    vis = [[False] * 8 for _ in range(8)]
    print(dfs(steps, vis, 0, 0, 0))

for level in LEVELS:
    print('-' * 60)
    solve(level)