codeingame/neighbors_sum_grids.py

'''
@author: olivier.massot, 2019
'''
import itertools
import sys
import time

def log(*msg):
    print(*msg, file=sys.stderr)
t0 = time.time()

# matrix = [[int(j) for j in input().split()] for _ in range(5)]
matrix = [[21, 14, 10, 0, 0], [0, 4, 1, 6, 0], [0, 3, 5, 11, 13], [15, 0, 0, 0, 0], [23, 24, 0, 20, 25]]
# matrix = [[19, 15, 11, 0, 23], [0, 0, 0, 0, 13], [0, 0, 1, 0, 22], [0, 0, 0, 0, 16], [14, 12, 21, 0, 25]]
# matrix = [[23, 22, 8, 0, 18], [0, 0, 0, 5, 0], [0, 0, 12, 0, 0], [0, 10, 0, 11, 9], [0, 0, 0, 0, 0]]
# matrix = [[0, 12, 0, 0, 0], [0, 0, 17, 0, 7], [0, 4, 3, 0, 18], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
log(matrix)

grid = {(x, y): v for y, row in enumerate(matrix) for x, v in enumerate(row)}

crd = [(x, y) for x in range(5) for y in range(5)]
ocrd = sorted(crd, key=lambda c: c[0] + c[1])

def get(grid, x, y):
    if x < 0 or y < 0:
        raise IndexError
    return grid[y][x]

def get_neighbors(x, y, d=5):
    return [(x+dx, y+dy) for dx in range(-1, 2) for dy in range(-1,2) \
            if (dx or dy) and 0 <= x + dx < d and 0 <= y + dy < d ]

neighbors = {c: get_neighbors(*c) for c in crd}

log("\n".join([" ".join([str(j).rjust(2) for j in row]) for row in matrix]))

def validate(grid, unused, start):
    x0, y0 = start
    scan = [(x, y) for x in range(x0 - 1, x0 + 2) for y in range(y0 - 1, y0 + 2) if (x, y) in grid]

    for c in scan:
        v = grid[c]
            
        if v <= 2:
            continue
        
        neighs = [grid.get(n) for n in get_neighbors(*c, d=len(grid))]
        neighs = [x for x in neighs if x is not None]
        if 0 in neighs:
            return True
        
        if not any(n1 + n2 == v for n1, n2 in itertools.combinations(neighs, 2)):
            return False
        
    return True
        
def test_candidates(grid):
    
    try:
        x, y = next(c for c, v in grid.items() if v == 0)
    except StopIteration:
        return grid
    
    unused = [j for j in range(25,0,-1) if not j in grid.values()]
    
    for j in unused:
        new_grid = dict(grid)
        new_grid[(x, y)] = j
        
        valid = validate(new_grid, unused, (x, y))
        
        if valid:
            r = test_candidates(new_grid)
            if r:
                return r

dim = 5

new_grid = {c: v for c, v in grid.items() if c[0] in range(dim) and c[1] in range(dim)}
print(new_grid)
r = test_candidates(new_grid)
print(r)

m = [[r[(x, y)] for x in range(dim)] for y in range(dim)]
log("\n".join([" ".join([str(j).rjust(2) for j in row]) for row in m]))

log(time.time() - t0)
for row in m:
    print(" ".join(map(str, row)))