fivethirtyeight riddler for April 1st 2016

peter_susan_and_barack.py

import itertools
import numpy as np 
from collections import Counter


def person_says_no(possibilities):
    uniqueness = Counter([d.values()[0] 
                         for d in possibilities]).items()
    impossibilities = [x for (x, count) in uniqueness if count == 1]

    return np.array([False if d.values()[0] in impossibilities else True 
                     for d in possibilities])

def person_says_yes(possibilities):
    uniqueness = Counter([d.values()[0] 
                         for d in possibilities]).items()
    answer = [x for (x, count) in uniqueness if count == 1]

    return [d.keys() for d in possibilities if d.values() == answer][0][0]

def update_on_lhs_no(lhs_possibilities, rhs_possibilities):
    lhs_inclusions = person_says_no(lhs_possibilities)
    lhs_possibilities = lhs_possibilities[lhs_inclusions]
    rhs_possibilities = rhs_possibilities[lhs_inclusions]

    return lhs_possibilities, rhs_possibilities


numbers = list(range(1, 10))
pairs = list(itertools.product(numbers, numbers))
pairs = [(l, r) for (l, r) in pairs if l <= r]

peter_poss = np.array([{(l, r) : l * r} for (l, r) in pairs])
susan_poss = np.array([{(l, r) : l + r} for (l, r) in pairs])

# Iterative rounds of saying no
peter_poss, susan_poss = update_on_lhs_no(peter_poss, susan_poss)
susan_poss, peter_poss = update_on_lhs_no(susan_poss, peter_poss)

peter_poss, susan_poss = update_on_lhs_no(peter_poss, susan_poss)
susan_poss, peter_poss = update_on_lhs_no(susan_poss, peter_poss)

peter_poss, susan_poss = update_on_lhs_no(peter_poss, susan_poss)
susan_poss, peter_poss = update_on_lhs_no(susan_poss, peter_poss)

peter_poss, susan_poss = update_on_lhs_no(peter_poss, susan_poss)
susan_poss, peter_poss = update_on_lhs_no(susan_poss, peter_poss)

# The answer 
person_says_yes(peter_poss)