Yajilin solver with gurobipy

solver.py

import random
from collections import namedtuple
import time
import gurobipy as gp
from gurobipy import GRB
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import json

with open("gurobi.lic", "r") as f:
    env = gp.Env(params=json.load(f))
env.setParam("OutputFlag", 0)

def decode_yajilin(data, difficulty=None):
    parts = data.split("/")
    w, h = int(parts[0]), int(parts[1])
    arrows = []
    bstr = parts[2]
    c = 0
    i = 0
    while i < len(bstr):
        ca = bstr[i]
        if ca in "01234":
            qdir = int(ca, 16)
            if bstr[i+1] == ".":
                qnum = -2
            else:
                qnum = int(bstr[i+1], 16)
            i += 2
        elif ca in "56789":
            qdir = int(ca, 16) - 5
            qnum = int(bstr[i+1:i+3], 16)
            i += 3
        elif "a" <= ca <= "z":
            skip = int(ca, 36) - 10
            c += skip + 1
            i += 1
            continue
        else:
            raise ValueError()
        y, x = divmod(c, w)
        arrows.append((y, x, qnum, [1, 3, 2, 0][qdir - 1]))
        c += 1
    return Challenge(w, h, arrows, difficulty, "https://puzz.link/p?yajilin/" + data)


def plot_yajilin_cycle(walls, dirs, arrows, h, w):
    da = [1, 0, -1, 0, 1]
    fig, ax = plt.subplots(figsize=(w/2, h/2), dpi=300)
    ax.set_xlim(0, w)
    ax.set_ylim(0, h)
    ax.set_xticks(range(w + 1))
    ax.set_yticks(range(h + 1))
    ax.invert_yaxis()
    ax.set_aspect("equal")
    ax.tick_params(labelbottom=False, labelleft=False, length=0)

    for ai, aj, count, d in arrows:
        arrow_sym = "→↑←↓"[d]
        count_str = "?" if count == -2 else str(count)
        ax.text(aj+0.5, ai+0.5, f"{count_str}{arrow_sym}", va="center", ha="center")
    for i in range(h):
        for j in range(w):
            if walls[i][j]:
                ax.add_patch(Rectangle((j, i), 1, 1, facecolor="black"))
            if dirs[i][j] in range(4):
                ax.plot([j+0.5, j + da[dirs[i][j]]+0.5], [i+0.5, i + da[dirs[i][j] + 1]+0.5], color="green", linewidth=2)

    ax.grid()
    plt.show()


Challenge = namedtuple("Challenge", ["w", "h", "arrows", "difficulty", "url"])

chals = []
with open("challenges", "r") as f:
    for line in f:
        data, difficulty = line.strip().split(",")
        chals.append(decode_yajilin(data, difficulty=difficulty))


def solve_MTZ_with_s(w, h, arrows, s):
    def oob(y, x):
        return not y in range(0, h) or not x in range(0, w)
    def ray(y, x, dir):
        while not oob(y, x):
            yield y, x
            x += da[dir]
            y += da[dir + 1]

    da = [1, 0, -1, 0, 1]
    model = gp.Model(env=env)

    wall_vars = model.addVars(h, w, vtype=GRB.BINARY, name="w")
    cycle_vars = model.addVars(h, w, vtype=GRB.INTEGER, name="c")
    edge_vars = model.addVars(h, w, 4, vtype=GRB.BINARY, name="e")

    arrow_pos = set((y, x) for y, x, _, _ in arrows)
    for y, x, c, dir in arrows:
        model.addConstr(wall_vars[y, x] == 0)
        model.addConstr(cycle_vars[y, x] == 0)
        model.addConstrs(edge_vars[y, x, k] == 0 for k in range(4))
        if c != -2:
            model.addConstr(gp.quicksum(wall_vars[i, j] for i, j in ray(y, x, dir)) == c)

    for i in range(h):
        for j in range(w):
            if (i, j) in arrow_pos:
                continue
            inedges = []
            outedges = []
            for d in range(4):
                ni, nj = i + da[d + 1], j + da[d]
                if oob(ni, nj) or (ni, nj) in arrow_pos:
                    model.addConstr(edge_vars[i, j, d] == 0)
                    continue
                inedges.append(edge_vars[ni, nj, (d + 2) % 4])
                outedges.append(edge_vars[i, j, d])
                if (ni, nj) != s:
                    model.addConstr(cycle_vars[i, j] - cycle_vars[ni, nj] + (h * w + 1) * edge_vars[i, j, d] <= h * w)
                model.addConstr(edge_vars[i, j, d] + wall_vars[i, j] + wall_vars[ni, nj] <= 1)
            model.addConstr(gp.quicksum(outedges) == 1 - wall_vars[i, j])
            model.addConstr(gp.quicksum(inedges) == 1 - wall_vars[i, j])
            model.addConstr(cycle_vars[i, j] <= h * w * (1 - wall_vars[i, j]))
            model.addConstr(1 - wall_vars[i, j] <= cycle_vars[i, j])

    model.setObjective(0, GRB.MINIMIZE)
    model.optimize()

    if model.Status == GRB.INFEASIBLE:
        return None

    walls = [[round(wall_vars[i, j].X) for j in range(w)] for i in range(h)]
    edges = [[[round(edge_vars[i, j, k].X) for k in range(4)] for j in range(w)] for i in range(h)]
    dirs = [[(edges[i][j] + [1]).index(1) for j in range(w)] for i in range(h)]

    return walls, dirs


def solve_MTZ(w, h, arrows):
    random.seed(0)
    arrow_pos = set((y, x) for y, x, _, _ in arrows)
    checked = set()
    while True:
        s = (random.randrange(0, h), random.randrange(0, w))
        if (s[0], s[1]) in arrow_pos or s in checked:
            continue
        result = solve_MTZ_with_s(w, h, arrows, s)
        checked.add(s)
        if result is not None:
            return result


def solve_lazy_elimination(w, h, arrows):
    def oob(y, x):
        return not y in range(0, h) or not x in range(0, w)
    def ray(y, x, dir):
        while not oob(y, x):
            yield y, x
            x += da[dir]
            y += da[dir + 1]

    da = [1, 0, -1, 0, 1]
    model = gp.Model(env=env)

    wall_vars = model.addVars(h, w, vtype=GRB.BINARY, name="w")
    edge_vars = model.addVars(h, w, 4, vtype=GRB.BINARY, name="e")

    arrow_pos = set((y, x) for y, x, _, _ in arrows)
    for y, x, c, dir in arrows:
        model.addConstr(wall_vars[y, x] == 0)
        model.addConstrs(edge_vars[y, x, k] == 0 for k in range(4))
        if c != -2:
            model.addConstr(gp.quicksum(wall_vars[i, j] for i, j in ray(y, x, dir)) == c)

    for i in range(h):
        for j in range(w):
            if (i, j) in arrow_pos:
                continue
            inedges = []
            outedges = []
            for d in range(4):
                ni, nj = i + da[d + 1], j + da[d]
                if oob(ni, nj) or (ni, nj) in arrow_pos:
                    model.addConstr(edge_vars[i, j, d] == 0)
                    continue
                inedges.append(edge_vars[ni, nj, (d + 2) % 4])
                outedges.append(edge_vars[i, j, d])
                model.addConstr(edge_vars[i, j, d] + wall_vars[i, j] + wall_vars[ni, nj] <= 1)
            model.addConstr(gp.quicksum(outedges) == 1 - wall_vars[i, j])
            model.addConstr(gp.quicksum(inedges) == 1 - wall_vars[i, j])

    def subtour_elim(model, where):
        if where == GRB.Callback.MIPSOL:
            wall_vals = model.cbGetSolution(wall_vars)
            edge_vals = model.cbGetSolution(edge_vars)
            num_walls = sum(wall_vals.values())
            visited = set()
            best_cycle = None
            for i in range(h):
                for j in range(w):
                    if (i, j) in visited or not any(round(edge_vals[i, j, d]) for d in range(4)):
                        continue
                    ci, cj = i, j
                    cycle = []
                    while True:
                        dir = [round(edge_vals[ci, cj, d]) for d in range(4)].index(1)
                        ni, nj = ci + da[dir + 1], cj + da[dir]
                        cycle.append(edge_vars[ci, cj, dir])
                        visited.add((ci, cj))
                        if (ni, nj) in visited:
                            break
                        ci, cj = ni, nj
                    if len(cycle) < h * w - len(arrows) - num_walls:
                        if best_cycle is None or len(best_cycle) < len(cycle):
                            best_cycle = cycle
            if best_cycle is not None:
                model.cbLazy(gp.quicksum(best_cycle) <= len(best_cycle) - 1)

    model.Params.LazyConstraints = 1
    model.Params.TimeLimit = 60
    model.setObjective(0, GRB.MINIMIZE)
    model.optimize(subtour_elim)

    if model.Status == GRB.Status.TIME_LIMIT:
        return None

    walls = [[round(wall_vars[i, j].X) for j in range(w)] for i in range(h)]
    edges = [[[round(edge_vars[i, j, k].X) for k in range(4)] for j in range(w)] for i in range(h)]
    dirs = [[(edges[i][j] + [1]).index(1) for j in range(w)] for i in range(h)]

    return walls, dirs


total = time.time()
for i, chal in enumerate(chals[::-1]):
    print(chal.difficulty, chal.url)
    random.seed(0)
    plot_yajilin_cycle([[0] * chal.w for _ in range(chal.h)], [[4] * chal.w for _ in range(chal.h)], chal.arrows, chal.h, chal.w)
    start = time.time()
    answer = solve_MTZ(chal.w, chal.h, chal.arrows)
    # answer = solve_lazy_elimination(chal.w, chal.h, chal.arrows)
    print(f"{i}: {time.time() - start:.3f} sec / {time.time() - total:.3f} sec")
    if answer is not None:
        walls, dirs = answer
        plot_yajilin_cycle(walls, dirs, chal.arrows, chal.h, chal.w)