https://www.reddit.com/r/dailyprogrammer/comments/4cw095/ challenge #260

neverendingsnake.py

from heapq import heappush, heappop
from numpy import array, copy

S = 's'
P = 'p'
F = '*'
O = 'O'
C = ' '


class Cell:
    def __init__(self, x, y, t):
        self.type = t
        self.f = float('inf')
        self.g = float('inf')
        self.x = x
        self.y = y

    def dist(self, other):
        if other is self:
            return float('inf')
        else:
            return abs(self.x - other.x) + abs(self.y - other.y)

    def __lt__(self, other):
        return self.f < other.f

    def __str__(self):
        return self.type


class Node:
    def __init__(self, matrix, bound, index, path):
        self.matrix = matrix
        self.bound = bound
        self.index = index
        self.path = path
        self.path.append(self)

    def __lt__(self, other):
        return self.bound < other.bound


def read_map(inp):
    locations = []
    goals = []
    start = None
    for i, line in enumerate(inp.split('\n')):
        if line.startswith('|'):
            locations.append([])
            for char in line:
                if char != '|':
                    if char == S:
                        start = Cell(len(locations[i - 1]), i - 1, S)
                        locations[i - 1].append(start)
                    elif char == F:
                        goal = Cell(len(locations[i - 1]), i - 1, F)
                        goals.append(goal)
                        locations[i - 1].append(goal)
                    else:
                        locations[i - 1].append(Cell(len(locations[i - 1]), i - 1, char))

    return locations, goals, start


def search(locations, goal, start):
    dirs = [(0, 1), (0, -1), (1, 0), (-1, 0)]
    closed = []
    heap = []
    start.f = start.dist(goal)
    start.g = 0
    start.type = P
    heappush(heap, start)

    while heap:
        q = heappop(heap)
        closed.append(q)
        if q is goal:
            break

        nbs = [locations[q.y + d[1]][q.x + d[0]] for d in dirs
               if 0 <= q.y + d[1] < len(locations) and 0 <= q.x + d[0] < len(locations[q.y + d[1]])]
        for n in nbs:
            if n in closed or n.type in (O, P):
                continue

            tf = q.g + 1
            if n not in heap:
                heappush(heap, n)
            elif tf >= n.g:
                continue

            n.p = q
            n.g = tf
            n.f = n.g + n.dist(goal)
    else:
        return False

    goal.type = S
    cur = goal.p
    while hasattr(cur, 'p'):
        cur.type = P
        cur = cur.p

    return True


def to_map(locations):
    result = '+%s+\n' % ''.ljust(len(locations[0]), '-')
    for row in locations:
        result += '|%s|\n' % ''.join(map(str, row))
    result += '+%s+' % ''.ljust(len(locations[0]), '-')

    return result


def reduce_matrix(matrix):
    total = 0
    for row in matrix:
        low = min(row)
        if low != float('inf'):
            total += low
            for i in xrange(len(row)):
                row[i] = row[i] - low

    for i in xrange(len(matrix)):
        col = matrix[:, i]
        low = min(col)
        if low != float('inf'):
            total += low
            for n in xrange(len(col)):
                col[n] = col[n] - low

    return total


def check_bound(matrix, s, d):
    total = matrix[s][d]
    matrix[d][s] = float('inf')
    matrix[s] = [float('inf') for _ in matrix[s]]
    matrix[:, d] = [float('inf') for _ in matrix[:, d]]

    return total


def try_path(m, start, path):
    cs, _, s = read_map(m)
    for g in path:
        if search(cs, cs[g.y][g.x], s):
            cs, _, s = read_map(to_map(cs))
        else:
            break
    else:
        s.type = 'e'
        cs[start.y][start.x] = 's'
        return to_map(cs)

    return ''


def branch_bound(m):
    cs, gs, s = read_map(m)
    gs.insert(0, s)

    l = []
    for g in gs:
        l.append([o.dist(g) for o in gs])

    nodes = []
    matrix = array(l)
    r = reduce_matrix(matrix)
    heappush(nodes, Node(matrix, r, 0, []))

    while True:
        cur = heappop(nodes)
        for i in xrange(len(cur.matrix)):
            if cur.matrix[cur.index][i] != float('inf') and i != 0:
                cp = copy(cur.matrix)
                total = cur.bound
                total += check_bound(cp, cur.index, i)
                node = Node(cp, total, i, cur.path[:])
                heappush(nodes, node)

                if len(node.path) == len(gs):
                    path = try_path(m, s, [gs[n] for n in [n.index for n in node.path] if n != 0])
                    if path:
                        return path