prbuddies.py

from z3 import *
import itertools

s = Solver()

names = ['Ali', 'Kelsey', 'Yuan', 'Chris', 'Narges', 'Weija', 'Jay', 'Linh', 'Joseph', 'Lucy']
nmap = {} # Maps string name to integer value
imap = {} # Reverse map, integer value to string name
for i, name in enumerate(names):
  nmap[name] = i
  imap[i] = name

# Create Z3 consts for each name
people = Ints(names)

for each in range(len(people)):
  # No one should ever be assigned to themselves
  s.add(people[each] != each)

  # Ensure everyone is assigned an integer we care about
  s.add(people[each] >= 0)
  s.add(people[each] < len(people))

# No two people should be assigned the same PR buddy at once
s.add(Distinct(people))

for person1, person2 in itertools.combinations(range(len(people)), 2):
  # If person1 is assigned to person2, it implies that person2 is assigned to person1
  s.add(Implies(people[person1] == person2, people[person2] == person1))

def block_existing(name, existing_buddies):
  for existing in existing_buddies:
    s.add(people[nmap[name]] != nmap[existing])

# These are the buddies we've already had
block_existing('Ali', ['Kelsey', 'Weija'])
block_existing('Kelsey', ['Ali', 'Linh'])
block_existing('Yuan', ['Jay', 'Joseph'])
block_existing('Chris', ['Joseph', 'Jay'])
block_existing('Narges', ['Weija', 'Ali'])
block_existing('Weija', ['Narges', 'Ali'])
block_existing('Jay', ['Yuan', 'Chris'])
block_existing('Linh', ['Narges', 'Kelsey'])
block_existing('Joseph', ['Chris', 'Yuan'])
block_existing('Lucy', ['Linh', 'Narges'])


def pp(model, decl):
  # We have to .as_long() because model[decl] returns an `IntNumRef`
  buddy_int_value = model[decl].as_long()
  buddy_name = imap[buddy_int_value]
  print(f"{decl}: {buddy_name}")

while s.check() == sat:
  print("-----------")
  m = s.model()
  for decl in m.decls():
   pp(m, decl)
  block = []
  for d in m:
    c = d()
    block.append(c != m[d])
  # And() together because we want *none* of the existing
  # assignments to ever happen again. Model blocks are usually
  # ORed, but that would allow people to be assigned the same
  # PR buddy later
  s.add(And(block))