advent-of-code-2019/day15.py

#!/usr/bin/env python3

import intcode

NORTH = 1
SOUTH = 2
WEST  = 3
EAST  = 4

DIRS = [NORTH, SOUTH, EAST, WEST]

OPPOSITE = {NORTH : SOUTH,
            SOUTH : NORTH,
            EAST  : WEST,
            WEST  : EAST
           }

WALL  = 0
EMPTY = 1
OXY   = 2

ENCODING = {WALL  : '#',
            EMPTY : '.',
            OXY   : 'O'
           }

def move(pos, d):
    if d == NORTH:
        return pos[0], pos[1] + 1
    if d == SOUTH:
        return pos[0], pos[1] - 1
    if d == EAST:
        return pos[0] + 1, pos[1]
    if d == WEST:
        return pos[0] - 1, pos[1]
    return NotImplemented

class Droid:
    def __init__(self, program):
        self.pos = 0, 0
        self.comp = intcode.emulator(program)

    def move(self, d):
        self.comp.write_input(d)
        out = next(self.comp) 
        if out != WALL:
            self.pos = move(self.pos, d)
        return out 

class Map(dict):
    def __init__(self, encoding):
        super().__init__()
        self.encoding = encoding

    def __str__(self):
        xs, ys = tuple(zip(*self.keys()))
        r = str()
        for y in range(min(ys), max(ys)+1):
            for x in range(min(xs), max(xs)+1):
                if x == y == 0:
                    r += 'X'
                elif (x, y) not in self:
                    r += ' '
                else:
                    r += self.encoding[self[(x, y)]]
            r += '\n'
        return r

def map_space(droid, space_map=None):
    if space_map == None:
        space_map = Map(ENCODING)
    for d in DIRS:
        new_pos = move(droid.pos, d)
        if new_pos in space_map:
            continue
        obj = droid.move(d)
        space_map[new_pos] = obj
        if obj == WALL:
            continue
        map_space(droid, space_map)
        droid.move(OPPOSITE[d])
    return space_map

def shortest_paths(space_map, pos, visited=None):
    if visited == None:
        visited = {pos : 0}
    path = visited[pos] + 1
    for d in DIRS:
        new = move(pos, d)
        if space_map[new] == WALL or (new in visited and visited[new] <= path):
            continue
        visited[new] = path
        shortest_paths(space_map, new, visited)
    return visited

def preproc(puzzle_input):
    program = intcode.parse(puzzle_input) 
    droid = Droid(program)
    space_map = map_space(droid)
    x, y = tuple(filter(lambda x: x[1]==OXY, space_map.items()))[0][0]
    path_map = shortest_paths(space_map, (x, y))
    return path_map

def partI(path_map):
    return path_map[(0, 0)]

def partII(path_map):
    return max(path_map.values())