algobots_set_cover.py

import json

data = json.load(open("../allData.json"))
traits_by_token = {item["tokenID"]: sorted(set(item["traits"])) for item in data}
assert len(data) == len(traits_by_token), "Duplicate tokens"
tokens = sorted(traits_by_token)
token_indices = {token: i for (i, token) in enumerate(tokens)}

traits = sorted({t for item in data for t in item["traits"]})
trait_indices = {trait: i for (i, trait) in enumerate(traits)}
tokens_by_trait = {trait: [] for trait in traits}
for (token, traits_) in traits_by_token.items():
    for trait in traits_:
        tokens_by_trait[trait].append(token)

num_allocated_variables = 0
def gensym():
    global num_allocated_variables
    # pre-increment; DIMACS variables are one-indexed
    num_allocated_variables += 1
    return num_allocated_variables

var_token_chosen = {token: gensym() for token in tokens}

TEAM_SIZE = 13

body_lines = []
def add_constraint(literals):
    body_lines.append(" ".join(str(x) for x in list(literals) + [0]))
def add_comment(comment):
    assert "\n" not in comment, "Bad comment"
    body_lines.append("c " + comment)
def add_section(name):
    add_comment("")
    add_comment(name)
    add_comment("---")

word_length = len(tokens).bit_length()
# binary encoding of team members
var_bit = {}
for team_member in range(TEAM_SIZE):
    for bit in range(word_length):
        var_bit[(team_member, bit)] = gensym()

add_section("one-hot encoding of team members")
var_team_relation = {}
for team_member in range(TEAM_SIZE):
    for token in tokens:
        var = gensym()
        var_team_relation[(team_member, token)] = var
        add_comment("team member #%d = token %s?" % (team_member, token))
        bits = bin(token_indices[token])[len("0b"):]
        bits = "0" * (word_length - len(bits)) + bits
        for (bit, value) in enumerate(bits):
            add_constraint([-var, {"1": 1, "0": -1}[value] * var_bit[(team_member, bit)]])

add_section("which tokens were chosen?")
for token in tokens:
    add_comment("was token %s chosen?" % token)
    constraint = [-var_token_chosen[token]]
    for team_member in range(TEAM_SIZE):
        constraint.append(var_team_relation[(team_member, token)])
    add_constraint(constraint)

add_section("all traits must be satisfied")
for trait in traits:
    add_comment("trait %r" % trait)
    add_constraint([var_token_chosen[token] for token in tokens_by_trait[trait]])

num_constraints = sum(not line.startswith("c ") for line in body_lines)
body_lines.insert(0, "p cnf %s %s" % (num_allocated_variables, num_constraints))
for line in body_lines:
    print(line)