EscapeTheJail

main.py

# Quick and dirty solution to the problem; code should be cleaned up
from pydtmc import MarkovChain


def solve_jail_escape(input_str):
    input_ = [list(k) for k in input_str.split('\n')]

    start_i, start_j = None, None
    target_i, target_j = None, None
    for i, row in enumerate(input_):
        for j, col in enumerate(row):
            if input_[i][j] == "@":
                input_[i][j] = "."
                start_i, start_j = i, j
            elif input_[i][j] == "$":
                input_[i][j] = "."
                target_i, target_j = i, j

    # MarkovChain gets confused by input "@..$#.." due to unreachable paths
    # The line below converts "@..$#.." to "@..$###"
    input_ = Graph(input_).select_island(start_i, start_j)

    ROWS = len(input_)
    COLS = len(input_[0])
    NODES = ROWS * COLS

    def conv(i, j) -> int:
        return i * COLS + j

    dx = [0, 1, 0, -1]
    dy = [1, 0, -1, 0]
    max_avail_moves = len(dx)

    markov_chain = []
    for i, row in enumerate(input_):
        for j, col in enumerate(row):
            contents = input_[i][j]
            markov_chain_row = [0] * NODES
            if contents == "#":
                markov_chain.append([0] * NODES)
                continue
            available_moves = max_avail_moves
            for i_node, (dx_, dy_) in enumerate(list(zip(dx, dy))):
                new_row = i + dx_
                new_col = j + dy_
                if not (0 <= new_row < ROWS):
                    available_moves = available_moves - 1
                    continue
                if not (0 <= new_col < COLS):
                    available_moves = available_moves - 1
                    continue
                new_contents = input_[new_row][new_col]
                if new_contents == "#":
                    available_moves = available_moves - 1
                    continue
                markov_chain_row[conv(new_row, new_col)] = 1
            if available_moves != 0:
                for k in range(len(markov_chain_row)):
                    markov_chain_row[k] = markov_chain_row[k] * 1 / available_moves
            markov_chain.append(markov_chain_row)

    target_k = conv(target_i, target_j)
    markov_chain[target_k] = [0] * NODES
    markov_chain[target_k][target_k] = 1

    for i, row in enumerate(markov_chain):
        if sum(row) == 0:
            new_row = [0] * NODES
            new_row[i] = 1
            markov_chain[i] = new_row

    names = [str(k + 1) for k in range(NODES)]
    mc = MarkovChain(markov_chain, names)

    mat = mc.mean_absorption_times()
    if mat is None:
        return -1

    answer_array = list(mat)
    start_k = conv(start_i, start_j)
    for i, row in enumerate(markov_chain):
        if row[i] == 1:
            answer_array.insert(i, 0)
    return answer_array[start_k]


class Graph:

    def __init__(self, g):
        ROW = len(g)
        COL = len(g[0])
        self.input_graph = g
        self.ROW = len(g)
        self.COL = len(g[0])
        new_g = [[(1 if g[i][j] in ('.', '$', '@') else 0) for j in range(COL)] for i in range(ROW)]
        self.graph = new_g

    # A function to check if a given cell
    # (row, col) can be included in DFS
    def isSafe(self, i, j, visited):
        # row number is in range, column number
        # is in range and value is 1
        # and not yet visited
        return (i >= 0 and i < self.ROW and
                j >= 0 and j < self.COL and
                not visited[i][j] and self.graph[i][j])

    # A utility function to do DFS for a 2D
    # boolean matrix. It only considers
    # the 8 neighbours as adjacent vertices
    def DFS(self, i, j, visited):

        # These arrays are used to get row and
        # column numbers of 8 neighbours
        # of a given cell
        rowNbr = [-1, 1, 0, 0]
        colNbr = [0, 0, -1, 1]

        # Mark this cell as visited
        visited[i][j] = True

        # Recur for all connected neighbours
        for k in range(8):
            if self.isSafe(i + rowNbr[k], j + colNbr[k], visited):
                self.DFS(i + rowNbr[k], j + colNbr[k], visited)

    # The main function that returns
    # count of islands in a given boolean
    # 2D matrix
    def countIslands(self):
        # Make a bool array to mark visited cells.
        # Initially all cells are unvisited
        visited = [[False for j in range(self.COL)] for i in range(self.ROW)]

        # Initialize count as 0 and traverse
        # through the all cells of
        # given matrix
        count = 0
        for i in range(self.ROW):
            for j in range(self.COL):
                # If a cell with value 1 is not visited yet,
                # then new island found
                if visited[i][j] == False and self.graph[i][j] == 1:
                    # Visit all cells in this island
                    # and increment island count
                    self.DFS(i, j, visited)
                    count += 1

        return count

    def select_island(self, start_i, start_j):
        visited = [[False for j in range(self.COL)] for i in range(self.ROW)]
        # Mark the selected island
        self.DFS(start_i, start_j, visited)
        output_graph = [["#" for j in range(self.COL)] for i in range(self.ROW)]
        for i in range(self.ROW):
            for j in range(self.COL):
                if visited[i][j]:
                    output_graph[i][j] = self.input_graph[i][j]
        return output_graph


assert solve_jail_escape("@$") == 1
assert solve_jail_escape("@..$") == 9
assert solve_jail_escape("$.\n.@") == 4
assert solve_jail_escape("@#\n#$") == -1
assert solve_jail_escape("@..$") == solve_jail_escape("..##@..$#...###")
print("Success")