from pulp import *
def chk(ary):
nh, nw = len(ary), len(ary[0])
mx, rw, rh, r4 = max(nw, nh), range(nw), range(nh), range(4)
m = LpProblem()
vb = [[LpVariable('b%d_%d' % (i, j), cat=LpBinary) for j in rw] for i in rh] # 0: red, 1: blue
# 0: left, 1: up, 2: right, 3:down
vd = [[[LpVariable('d%d_%d_%d' % (i, j, k), lowBound=0) for k in r4] for j in rw] for i in rh]
m += lpSum(vb[i][j] for i in rh for j in rw)
for i in rh:
for j in rw:
for k in r4:
m += vd[i][j][k] <= mx * vb[i][j]
ik, jk = i + [-1, 0, 1, 0][k], j + [0, -1, 0, 1][k]
if 0 <= ik < nh and 0 <= jk < nw:
m += vd[i][j][k] >= vd[ik][jk][k] + 1 - mx * (1 - vb[i][j])
m += vd[i][j][k] <= vd[ik][jk][k] + 1 + mx * (1 - vb[i][j])
else:
m += vd[i][j][k] == vb[i][j]
if ary[i][j] < 0:
m += vb[i][j] == 0
elif ary[i][j]:
m += vb[i][j] == 1
m += lpSum(vd[i][j]) == ary[i][j] + 4
m.solve()
if m.status != 1:
return None
res = [[int(value(vb[i][j])) for j in rw] for i in rh]
ll = [vb[i][j] for i in rh for j in rw if value(vb[i][j]) > 0.5]
m += lpSum(ll) <= len(ll) - 1
m.solve()
return res if m.status != 1 else None
chk([
[0,0,-1,3,0],
[0,0,0,0,0],
[0,0,0,0,5],
[3,4,0,0,-1],
[0,0,4,-1,0],
])
[[0, 0, 0, 1, 1], [0, 0, 0, 1, 1], [0, 1, 1, 1, 1], [1, 1, 1, 0, 0], [1, 1, 1, 0, 0]]
chk([
[-1,0,0,3,0],
[0,2,0,0,-1],
[4,4,0,0,0],
[5,0,5,0,0],
[0,0,0,3,0],
])
[[0, 0, 1, 1, 1], [0, 1, 0, 1, 0], [1, 1, 1, 0, 0], [1, 1, 1, 1, 0], [1, 0, 1, 1, 1]]