minesweeper.py

import argparse
import collections

# Let's define a global dispatch table for the 8 different possible neighbors
# clockwise starting top-left: tl, t, tr, r, br, b, bl, l.
# The target index only depends on the number of rows and the start index jj
# As we won't know beforehand which functions to use for a given start index,
# we make them all dependent on both jj and n

neighbordispatch = {
    'tl': lambda jj, n: jj - n - 1,
    't': lambda jj, n: jj - n,
    'tr': lambda jj, n: jj - n + 1,
    'r': lambda jj, n: jj + 1,
    'br': lambda jj, n: jj + n + 1,
    'b': lambda jj, n: jj + n,
    'bl': lambda jj, n: jj + n - 1,
    'l': lambda jj, n: jj - 1,
}

def getinput(teststr):
    """Cleans up an input line and splits it up"""
    idxs, mines = teststr.split(';')
    m, n = [int(item) for item in idxs.split(',')]
    return m, n, mines
    
def neighbors(jj, m, n):
    """Generates set of valid neighbors for a given index"""
    neighbors = set(neighbordispatch.keys())
    if jj < n:
        # first row, take out the neighbor candidates above
        neighbors.difference_update({'tl', 't', 'tr'})
    if jj >= n * (m-1):
        # last row, take out the neighbor candidates below
        neighbors.difference_update({'bl', 'b', 'br'})
    if jj % n == 0:
        # first column, take out the neighbor candidates to the left
        neighbors.difference_update({'tl', 'l', 'bl'})
    if (jj+1) % n == 0:
        # last column, take out the neighbor candidates to the right
        neighbors.difference_update({'tr', 'r', 'br'})
    return neighbors    

def selectmines(arg1, arg2):
    return arg1 if arg1 == '*' else str(arg2)
    
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("infilepath")
    args = parser.parse_args()

    with open(args.infilepath, "rU") as source:
        for line in source:
            teststr = line.strip()
            if not teststr:
                continue
            m, n, mines = getinput(teststr)
            neighborcounts = [
                countr['*'] for countr in [
                    collections.Counter([
                        mines[neighbordispatch[neigh](ii, n)]
                        for neigh in neighbors(ii, m, n)
                    ]) for ii in range(m*n)
                ]
            ]
            result = map(selectmines, mines, neighborcounts)
            print(''.join(result))

if __name__ == '__main__':
    main()