In [21]:

    

sinput="""2
3 4 2
1 1 1
3 3 2
5 5 2
4 1 2
3 2 2"""

import logging
logging.basicConfig()
log = logging.getLogger('jam')

import numpy as np
import pandas as pd


    

In [15]:

    

def parse_input(inp):
    lines = [[int(n) for n in line.split()] for line in inp.split('\n')]
    casenum = lines.pop(0)[0]
    cases = []
    for _ in range(casenum):
        grid = lines.pop(0)
        S = grid[2]
        stations = []
        for _ in range(S):
            stations.append(lines.pop(0))
        cases.append((grid, stations))
    
    return cases
        
parse_input(sinput)


    

    
Out[15]:





[([3, 4, 2], [[1, 1, 1], [3, 3, 2]]), ([5, 5, 2], [[4, 1, 2], [3, 2, 2]])]

In [46]:

    

class Grid:
    
    def __init__(self, stations):
        self.stations = stations
        
        
class Station:
    
    def __init__(self):
        self.size = 0
        self.weighted_size = 0
        self.block = 0


    

In [45]:

    

def assign_grid(case):
    shape = case[0]
    r, c = shape[0], shape[1]
    grid = [set() for _ in range(r*c)]
    
    for x, station in enumerate(case[1]):
        sn = x+1
        reach = station[2]
        sr, sc = station[0], station[1]
        grid[(sr-1)*c + sc - 1] = sn
        for i in range(max(1, sr-reach), min(r, sr+reach)+1):
            for j in range(max(1, sc-reach), min(c, sc+reach)+1):
                try:
                    x = (i-1)*c + j - 1
                    grid[x].add(sn)
                except AttributeError:
                    pass
                    
    return grid
    
cases = parse_input(sinput)
maps = [assign_grid(c) for c in cases]
arr = np.array(maps[1])
arr.reshape(5,5)


    

    
Out[45]:





array([[{2}, {2}, {2}, {2}, set()],
       [{1, 2}, {1, 2}, {1, 2}, {2}, set()],
       [{1, 2}, 2, {1, 2}, {2}, set()],
       [1, {1, 2}, {1, 2}, {2}, set()],
       [{1, 2}, {1, 2}, {1, 2}, {2}, set()]], dtype=object)

In [ ]:

    

def