Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment